From cfac25c0fb90851aa79acee42f03720e3ba42bc6 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Wed, 17 Nov 2021 18:56:56 +0800 Subject: [PATCH 001/128] Add support union type Signed-off-by: Kevin Su Signed-off-by: maximsmol --- .github/workflows/pythonbuild.yml | 2 + flytekit/core/type_engine.py | 57 ++++++++++++++++++-- flytekit/models/types.py | 32 +++++++++++ tests/flytekit/unit/core/test_type_engine.py | 33 ++++++++++++ tests/flytekit/unit/core/test_type_hints.py | 40 ++++++++++++++ 5 files changed, 160 insertions(+), 4 deletions(-) diff --git a/.github/workflows/pythonbuild.yml b/.github/workflows/pythonbuild.yml index 924e33fa63..b3fb526395 100644 --- a/.github/workflows/pythonbuild.yml +++ b/.github/workflows/pythonbuild.yml @@ -36,6 +36,7 @@ jobs: run: | python -m pip install --upgrade pip==21.2.4 setuptools wheel make setup${{ matrix.spark-version-suffix }} + git clone https://github.com/pingsutw/flyteidl.git && cd flyteidl && git checkout union_type && pip install . && cd .. pip freeze - name: Test with coverage run: | @@ -93,6 +94,7 @@ jobs: cd plugins/${{ matrix.plugin-names }} pip install -e . pip install --no-deps -U https://github.com/flyteorg/flytekit/archive/${{ github.sha }}.zip#egg=flytekit + git clone https://github.com/pingsutw/flyteidl.git && cd flyteidl && git checkout union_type && pip install . && cd .. pip freeze - name: Test with coverage run: | diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 925f79c770..59656832b4 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -28,7 +28,7 @@ from flytekit.models import types as _type_models from flytekit.models.core import types as _core_types from flytekit.models.literals import Literal, LiteralCollection, LiteralMap, Primitive, Scalar -from flytekit.models.types import LiteralType, SimpleType +from flytekit.models.types import LiteralType, SimpleType, UnionType T = typing.TypeVar("T") DEFINITIONS = "definitions" @@ -444,8 +444,6 @@ def to_literal(cls, ctx: FlyteContext, python_val: typing.Any, python_type: Type """ Converts a python value of a given type and expected ``LiteralType`` into a resolved ``Literal`` value. """ - if python_val is None: - raise AssertionError(f"Python value cannot be None, expected {python_type}/{expected}") transformer = cls.get_transformer(python_type) if transformer.type_assertions_enabled: transformer.assert_type(python_type, python_val) @@ -596,6 +594,55 @@ def guess_python_type(self, literal_type: LiteralType) -> Type[list]: raise ValueError(f"List transformer cannot reverse {literal_type}") +class UnionTransformer(TypeTransformer[T]): + """ + Transformer that handles a univariate typing.Union[T] + """ + + def __init__(self): + super().__init__("Typed Union", typing.Union) + + @staticmethod + def get_sub_type(t: Type[T]) -> Type[T]: + """ + Return the generic Type T of the Union + """ + if hasattr(t, "__origin__") and t.__origin__ is typing.Union: # type: ignore + if hasattr(t, "__args__"): + return t.__args__ # type: ignore + raise ValueError("Only generic univariate typing.Union[T] type is supported.") + + def get_literal_type(self, t: Type[T]) -> Optional[LiteralType]: + try: + sub_type = [TypeEngine.to_literal_type(v) if v else "" for v in self.get_sub_type(t)] + return _type_models.LiteralType(union_type=UnionType(sub_type)) + except Exception as e: + raise ValueError(f"Type of Generic Union type is not supported, {e}") + + def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], expected: LiteralType) -> Literal: + t = type(python_val) + return TypeEngine.to_literal(ctx, python_val, t, expected) + + def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[T]) -> Optional[typing.Any]: + st = self.get_sub_type(expected_python_type) + has_none_type = False + for v in st: + try: + if isinstance(v, type(None)): + has_none_type = True + val = TypeEngine.to_python_value(ctx, lv, v) + if val: + return val + except Exception as e: + logger.debug(f"Failed to convert from {lv} to {v}") + if has_none_type: + return True + raise TypeError(f"Cannot convert from {lv} to {expected_python_type}") + + def guess_python_type(self, literal_type: LiteralType) -> type: + return TypeEngine.guess_python_type(literal_type) + + class DictTransformer(TypeTransformer[dict]): """ Transformer that transforms a univariate dictionary Dict[str, T] to a Literal Map or @@ -907,9 +954,11 @@ def _register_default_type_transformers(): _type_models.LiteralType(simple=_type_models.SimpleType.NONE), lambda x: None, lambda x: None, - ) + ), + [type(None)], ) TypeEngine.register(ListTransformer()) + TypeEngine.register(UnionTransformer()) TypeEngine.register(DictTransformer()) TypeEngine.register(TextIOTransformer()) TypeEngine.register(BinaryIOTransformer()) diff --git a/flytekit/models/types.py b/flytekit/models/types.py index 719cdfd7b0..664d466e0c 100644 --- a/flytekit/models/types.py +++ b/flytekit/models/types.py @@ -1,4 +1,5 @@ import json as _json +import typing from flyteidl.core import types_pb2 as _types_pb2 from google.protobuf import json_format as _json_format @@ -98,6 +99,28 @@ def from_flyte_idl(cls, proto): return cls(columns=[SchemaType.SchemaColumn.from_flyte_idl(c) for c in proto.columns]) +class UnionType(_common.FlyteIdlEntity): + """ + Models _types_pb2.UnionType + """ + + def __init__(self, values: typing.List["LiteralType"]): + self._values = values + + @property + def values(self) -> typing.List["LiteralType"]: + return self._values + + def to_flyte_idl(self) -> _types_pb2.UnionType: + return _types_pb2.UnionType( + values=[val.to_flyte_idl() if val else None for val in self._values], + ) + + @classmethod + def from_flyte_idl(cls, proto: _types_pb2.UnionType): + return cls(values=proto.values) + + class LiteralType(_common.FlyteIdlEntity): def __init__( self, @@ -107,6 +130,7 @@ def __init__( map_value_type=None, blob=None, enum_type=None, + union_type=None, metadata=None, ): """ @@ -119,6 +143,7 @@ def __init__( string. :param flytekit.models.core.types.BlobType blob: For blob objects, this describes the type. :param flytekit.models.core.types.EnumType enum_type: For enum objects, describes an enum + :param flytekit.models.core.types.UnionType union_type: For union objects, describes an python union type. :param dict[Text, T] metadata: Additional data describing the type """ self._simple = simple @@ -127,6 +152,7 @@ def __init__( self._map_value_type = map_value_type self._blob = blob self._enum_type = enum_type + self._union_type = union_type self._metadata = metadata @property @@ -159,6 +185,10 @@ def blob(self) -> _core_types.BlobType: def enum_type(self) -> _core_types.EnumType: return self._enum_type + @property + def union_type(self) -> UnionType: + return self._union_type + @property def metadata(self): """ @@ -185,6 +215,7 @@ def to_flyte_idl(self): map_value_type=self.map_value_type.to_flyte_idl() if self.map_value_type is not None else None, blob=self.blob.to_flyte_idl() if self.blob is not None else None, enum_type=self.enum_type.to_flyte_idl() if self.enum_type else None, + union_type=self.union_type.to_flyte_idl() if self.union_type else None, metadata=metadata, ) return t @@ -208,6 +239,7 @@ def from_flyte_idl(cls, proto): map_value_type=map_value_type, blob=_core_types.BlobType.from_flyte_idl(proto.blob) if proto.HasField("blob") else None, enum_type=_core_types.EnumType.from_flyte_idl(proto.enum_type) if proto.HasField("enum_type") else None, + union_type=UnionType.from_flyte_idl(proto.union_type) if proto.HasField("union_type") else None, metadata=_json_format.MessageToDict(proto.metadata) or None, ) diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 6c660b7106..1f77828373 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -547,6 +547,39 @@ def test_enum_type(): TypeEngine.to_literal_type(UnsupportedEnumValues) +def test_union_type(): + pt = typing.Union[str, int] + lt = TypeEngine.to_literal_type(pt) + assert lt.union_type.values == [LiteralType(simple=SimpleType.STRING), LiteralType(simple=SimpleType.INTEGER)] + + ctx = FlyteContextManager.current_context() + lv = TypeEngine.to_literal(ctx, 3, pt, lt) + v = TypeEngine.to_python_value(ctx, lv, pt) + assert lv.scalar.primitive.integer == 3 + assert v == 3 + + lv = TypeEngine.to_literal(ctx, "hello", pt, lt) + v = TypeEngine.to_python_value(ctx, lv, pt) + assert lv.scalar.primitive.string_value == "hello" + assert v == "hello" + + +def test_optional_type(): + pt = typing.Optional[int] + lt = TypeEngine.to_literal_type(pt) + assert lt.union_type.values == [LiteralType(simple=SimpleType.INTEGER), LiteralType(simple=SimpleType.NONE)] + + ctx = FlyteContextManager.current_context() + lv = TypeEngine.to_literal(ctx, 3, pt, lt) + v = TypeEngine.to_python_value(ctx, lv, pt) + assert lv.scalar.primitive.integer == 3 + assert v == 3 + + lv = TypeEngine.to_literal(ctx, None, pt, lt) + v = TypeEngine.to_python_value(ctx, lv, pt) + assert v is None + + @pytest.mark.parametrize( "python_value,python_types,expected_literal_map", [ diff --git a/tests/flytekit/unit/core/test_type_hints.py b/tests/flytekit/unit/core/test_type_hints.py index e9f54c829b..40f9d7796d 100644 --- a/tests/flytekit/unit/core/test_type_hints.py +++ b/tests/flytekit/unit/core/test_type_hints.py @@ -1427,3 +1427,43 @@ def foo3(a: typing.Dict) -> typing.Dict: with pytest.raises(TypeError, match="Not a collection type simple: STRUCT\n but got a list \\[{'hello': 2}\\]"): foo3(a=[{"hello": 2}]) + + +def test_union_type(): + @task + def t1(a: typing.Union[int, str]) -> typing.Union[int, str]: + return a + + @workflow + def wf(a: typing.Union[int, str]) -> typing.Union[int, str]: + return t1(a=a) + + assert wf(a=2) == 2 + assert wf(a="2") == "2" + + @task + def t1(a: typing.Union[float, dict]) -> typing.Union[float, dict]: + return a + + @workflow + def wf(a: typing.Union[int, str]) -> typing.Union[int, str]: + return t1(a=a) + + with pytest.raises( + TypeError, + match='Cannot convert from scalar {\n primitive {\n string_value: "2"\n }\n}\n to typing.Union\[float, dict\]', + ): + assert wf(a="2") == "2" + + +def test_optional_type(): + @task + def t1(a: typing.Optional[int]) -> typing.Optional[int]: + return a + + @workflow + def wf(a: typing.Optional[int]) -> typing.Optional[int]: + return t1(a=a) + + assert wf(a=2) == 2 + assert wf(a=None) is None From 1027a3137237ccfb45ea49a1583cbb5f1c474dee Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Thu, 18 Nov 2021 01:09:37 +0800 Subject: [PATCH 002/128] Fixed test Signed-off-by: Kevin Su Signed-off-by: maximsmol --- flytekit/core/type_engine.py | 13 ++++++------- tests/flytekit/unit/core/test_type_hints.py | 10 +++++----- 2 files changed, 11 insertions(+), 12 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 59656832b4..a0fae9c464 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -625,22 +625,21 @@ def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], exp def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[T]) -> Optional[typing.Any]: st = self.get_sub_type(expected_python_type) - has_none_type = False for v in st: try: - if isinstance(v, type(None)): - has_none_type = True + if v == type(None) and lv is None: + return None val = TypeEngine.to_python_value(ctx, lv, v) if val: return val except Exception as e: - logger.debug(f"Failed to convert from {lv} to {v}") - if has_none_type: - return True + logger.debug(f"Failed to convert from {lv} to {v}", e) raise TypeError(f"Cannot convert from {lv} to {expected_python_type}") def guess_python_type(self, literal_type: LiteralType) -> type: - return TypeEngine.guess_python_type(literal_type) + if literal_type.union_type: + return typing.Union[[TypeEngine.guess_python_type(v) for v in literal_type.union_type.values]] + raise ValueError(f"Union transformer cannot reverse {literal_type}") class DictTransformer(TypeTransformer[dict]): diff --git a/tests/flytekit/unit/core/test_type_hints.py b/tests/flytekit/unit/core/test_type_hints.py index 40f9d7796d..a3c626e1fa 100644 --- a/tests/flytekit/unit/core/test_type_hints.py +++ b/tests/flytekit/unit/core/test_type_hints.py @@ -1442,18 +1442,18 @@ def wf(a: typing.Union[int, str]) -> typing.Union[int, str]: assert wf(a="2") == "2" @task - def t1(a: typing.Union[float, dict]) -> typing.Union[float, dict]: + def t2(a: typing.Union[float, dict]) -> typing.Union[float, dict]: return a @workflow - def wf(a: typing.Union[int, str]) -> typing.Union[int, str]: - return t1(a=a) + def wf2(a: typing.Union[int, str]) -> typing.Union[int, str]: + return t2(a=a) with pytest.raises( TypeError, - match='Cannot convert from scalar {\n primitive {\n string_value: "2"\n }\n}\n to typing.Union\[float, dict\]', + match='Cannot convert from scalar {\n primitive {\n string_value: "2"\n }\n}\n to typing.Union\\[float, dict\\]', ): - assert wf(a="2") == "2" + assert wf2(a="2") == "2" def test_optional_type(): From eab9b28626f2c4a47d2392ab7862607c268e2f3c Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Thu, 18 Nov 2021 01:55:42 +0800 Subject: [PATCH 003/128] Fixed test Signed-off-by: Kevin Su Signed-off-by: maximsmol --- flytekit/core/type_engine.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index a0fae9c464..3116d11359 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -444,6 +444,8 @@ def to_literal(cls, ctx: FlyteContext, python_val: typing.Any, python_type: Type """ Converts a python value of a given type and expected ``LiteralType`` into a resolved ``Literal`` value. """ + if python_val is None and expected.union_type is None: + raise AssertionError(f"Python value cannot be None, expected {python_type}/{expected}") transformer = cls.get_transformer(python_type) if transformer.type_assertions_enabled: transformer.assert_type(python_type, python_val) From 4a0fdb9a4b9cbd0a73a3e10b5d3a01dd4995e994 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Thu, 18 Nov 2021 02:16:43 +0800 Subject: [PATCH 004/128] Fixed lint Signed-off-by: Kevin Su Signed-off-by: maximsmol --- flytekit/core/type_engine.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 3116d11359..cb16b9e3c0 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -626,11 +626,11 @@ def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], exp return TypeEngine.to_literal(ctx, python_val, t, expected) def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[T]) -> Optional[typing.Any]: + if lv is None: + return None st = self.get_sub_type(expected_python_type) for v in st: try: - if v == type(None) and lv is None: - return None val = TypeEngine.to_python_value(ctx, lv, v) if val: return val From 38a8f7eead0379136c5b5f79ffa55881ad5b5d1f Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Thu, 18 Nov 2021 17:33:53 +0800 Subject: [PATCH 005/128] Fixed tests Signed-off-by: Kevin Su Signed-off-by: maximsmol --- .github/workflows/pythonbuild.yml | 3 ++- flytekit/core/type_engine.py | 2 +- flytekit/models/types.py | 2 +- 3 files changed, 4 insertions(+), 3 deletions(-) diff --git a/.github/workflows/pythonbuild.yml b/.github/workflows/pythonbuild.yml index b3fb526395..9992db6f3e 100644 --- a/.github/workflows/pythonbuild.yml +++ b/.github/workflows/pythonbuild.yml @@ -36,7 +36,7 @@ jobs: run: | python -m pip install --upgrade pip==21.2.4 setuptools wheel make setup${{ matrix.spark-version-suffix }} - git clone https://github.com/pingsutw/flyteidl.git && cd flyteidl && git checkout union_type && pip install . && cd .. + git clone https://github.com/flyteorg/flyteidl.git && cd flyteidl && git checkout union_type && pip install . && cd .. pip freeze - name: Test with coverage run: | @@ -142,5 +142,6 @@ jobs: run: | python -m pip install --upgrade pip==21.2.4 setuptools wheel pip install -r doc-requirements.txt + git clone https://github.com/flyteorg/flyteidl.git && cd flyteidl && git checkout union_type && pip install . && cd .. - name: Build the documentation run: make -C docs html diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index cb16b9e3c0..963f0a9130 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -640,7 +640,7 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: def guess_python_type(self, literal_type: LiteralType) -> type: if literal_type.union_type: - return typing.Union[[TypeEngine.guess_python_type(v) for v in literal_type.union_type.values]] + return typing.Union[tuple(TypeEngine.guess_python_type(v) for v in literal_type.union_type.values)] raise ValueError(f"Union transformer cannot reverse {literal_type}") diff --git a/flytekit/models/types.py b/flytekit/models/types.py index 664d466e0c..7e05fd1593 100644 --- a/flytekit/models/types.py +++ b/flytekit/models/types.py @@ -118,7 +118,7 @@ def to_flyte_idl(self) -> _types_pb2.UnionType: @classmethod def from_flyte_idl(cls, proto: _types_pb2.UnionType): - return cls(values=proto.values) + return cls(values=[LiteralType.from_flyte_idl(v) for v in proto.values]) class LiteralType(_common.FlyteIdlEntity): From e179fb97af9733baf402894aba169e67da6791c3 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Thu, 18 Nov 2021 17:58:25 +0800 Subject: [PATCH 006/128] Updated tests Signed-off-by: Kevin Su Signed-off-by: maximsmol --- .github/workflows/pythonbuild.yml | 2 +- flytekit/core/type_engine.py | 8 ++++++-- tests/flytekit/unit/core/test_type_hints.py | 14 ++++++++++++-- 3 files changed, 19 insertions(+), 5 deletions(-) diff --git a/.github/workflows/pythonbuild.yml b/.github/workflows/pythonbuild.yml index 9992db6f3e..b83bd79bfd 100644 --- a/.github/workflows/pythonbuild.yml +++ b/.github/workflows/pythonbuild.yml @@ -94,7 +94,7 @@ jobs: cd plugins/${{ matrix.plugin-names }} pip install -e . pip install --no-deps -U https://github.com/flyteorg/flytekit/archive/${{ github.sha }}.zip#egg=flytekit - git clone https://github.com/pingsutw/flyteidl.git && cd flyteidl && git checkout union_type && pip install . && cd .. + git clone https://github.com/flyteorg/flyteidl.git && cd flyteidl && git checkout union_type && pip install . && cd .. pip freeze - name: Test with coverage run: | diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 963f0a9130..1b95637b68 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -622,8 +622,12 @@ def get_literal_type(self, t: Type[T]) -> Optional[LiteralType]: raise ValueError(f"Type of Generic Union type is not supported, {e}") def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], expected: LiteralType) -> Literal: - t = type(python_val) - return TypeEngine.to_literal(ctx, python_val, t, expected) + for t in python_type.__args__: + try: + return TypeEngine.to_literal(ctx, python_val, t, expected) + except Exception as e: + logger.debug(f"Failed to convert from {python_val} to {t}", e) + raise TypeError(f"Cannot convert from {python_val} to {python_type}") def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[T]) -> Optional[typing.Any]: if lv is None: diff --git a/tests/flytekit/unit/core/test_type_hints.py b/tests/flytekit/unit/core/test_type_hints.py index a3c626e1fa..05edf43d6e 100644 --- a/tests/flytekit/unit/core/test_type_hints.py +++ b/tests/flytekit/unit/core/test_type_hints.py @@ -3,6 +3,7 @@ import functools import os import random +import tempfile import typing from collections import OrderedDict from dataclasses import dataclass @@ -31,6 +32,7 @@ from flytekit.models.interface import Parameter from flytekit.models.task import Resources as _resource_models from flytekit.models.types import LiteralType, SimpleType +from flytekit.types.file import FlyteFile from flytekit.types.schema import FlyteSchema, SchemaOpenMode serialization_settings = context_manager.SerializationSettings( @@ -1430,16 +1432,24 @@ def foo3(a: typing.Dict) -> typing.Dict: def test_union_type(): + ut = typing.Union[int, str, float, FlyteFile, FlyteSchema, typing.List[int], typing.Dict[str, int]] + @task - def t1(a: typing.Union[int, str]) -> typing.Union[int, str]: + def t1(a: ut) -> ut: return a @workflow - def wf(a: typing.Union[int, str]) -> typing.Union[int, str]: + def wf(a: ut) -> ut: return t1(a=a) assert wf(a=2) == 2 assert wf(a="2") == "2" + assert wf(a=2.0) == 2.0 + file = tempfile.NamedTemporaryFile() + assert isinstance(wf(a=FlyteFile(file.name)), FlyteFile) + assert isinstance(wf(a=FlyteSchema()), FlyteSchema) + assert wf(a=[1, 2, 3]) == [1, 2, 3] + assert wf(a={"a": 1}) == {"a": 1} @task def t2(a: typing.Union[float, dict]) -> typing.Union[float, dict]: From c1f8cfebea1242307b7ac30464e19ea915079525 Mon Sep 17 00:00:00 2001 From: Maksim Smolin Date: Fri, 3 Dec 2021 09:27:43 -0800 Subject: [PATCH 007/128] Update Union to use tagged Unions Signed-off-by: maximsmol --- flytekit/core/type_engine.py | 45 ++++++++++++++------------ flytekit/models/literals.py | 62 +++++++++++++++++++++++++++++++++++- flytekit/models/types.py | 12 +++---- 3 files changed, 91 insertions(+), 28 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 1b95637b68..0fa60808b9 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -27,7 +27,7 @@ from flytekit.models import interface as _interface_models from flytekit.models import types as _type_models from flytekit.models.core import types as _core_types -from flytekit.models.literals import Literal, LiteralCollection, LiteralMap, Primitive, Scalar +from flytekit.models.literals import Literal, LiteralCollection, LiteralMap, Primitive, Scalar, Union from flytekit.models.types import LiteralType, SimpleType, UnionType T = typing.TypeVar("T") @@ -604,47 +604,50 @@ class UnionTransformer(TypeTransformer[T]): def __init__(self): super().__init__("Typed Union", typing.Union) - @staticmethod - def get_sub_type(t: Type[T]) -> Type[T]: - """ - Return the generic Type T of the Union - """ - if hasattr(t, "__origin__") and t.__origin__ is typing.Union: # type: ignore - if hasattr(t, "__args__"): - return t.__args__ # type: ignore - raise ValueError("Only generic univariate typing.Union[T] type is supported.") - def get_literal_type(self, t: Type[T]) -> Optional[LiteralType]: try: - sub_type = [TypeEngine.to_literal_type(v) if v else "" for v in self.get_sub_type(t)] + sub_type = [TypeEngine.to_literal_type(v) for v in typing.get_args(t)] return _type_models.LiteralType(union_type=UnionType(sub_type)) except Exception as e: raise ValueError(f"Type of Generic Union type is not supported, {e}") def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], expected: LiteralType) -> Literal: - for t in python_type.__args__: + for idx, t in enumerate(typing.get_args(python_type)): try: - return TypeEngine.to_literal(ctx, python_val, t, expected) + val = TypeEngine.to_literal(ctx, python_val, t, expected) except Exception as e: logger.debug(f"Failed to convert from {python_val} to {t}", e) + continue + + return Literal(scalar=Scalar(union=Union(value=val, type=expected, tag=idx))) + raise TypeError(f"Cannot convert from {python_val} to {python_type}") def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[T]) -> Optional[typing.Any]: - if lv is None: - return None - st = self.get_sub_type(expected_python_type) - for v in st: + for x in typing.get_args(expected_python_type): + if x is type(None): + if lv is None: # todo(maximsmol): why would this ever be None and not Void? + return None + + if lv.scalar is not None and lv.scalar.none_type is not None: + return None + + break + + for v in typing.get_args(expected_python_type): try: val = TypeEngine.to_python_value(ctx, lv, v) if val: return val - except Exception as e: + except Exception as e: # todo(maximsmol): separate programmer errors from type conversion failures logger.debug(f"Failed to convert from {lv} to {v}", e) + raise TypeError(f"Cannot convert from {lv} to {expected_python_type}") def guess_python_type(self, literal_type: LiteralType) -> type: - if literal_type.union_type: - return typing.Union[tuple(TypeEngine.guess_python_type(v) for v in literal_type.union_type.values)] + if literal_type.union_type is not None: + return typing.Union[tuple(TypeEngine.guess_python_type(v) for v in literal_type.union_type.variants)] + raise ValueError(f"Union transformer cannot reverse {literal_type}") diff --git a/flytekit/models/literals.py b/flytekit/models/literals.py index 48553005b8..e45263eca3 100644 --- a/flytekit/models/literals.py +++ b/flytekit/models/literals.py @@ -9,6 +9,7 @@ from flytekit.models.core import types as _core_types from flytekit.models.types import OutputReference as _OutputReference from flytekit.models.types import SchemaType as _SchemaType +from flytekit.models.types import UnionType as _UnionType class RetryStrategy(_common.FlyteIdlEntity): @@ -546,6 +547,54 @@ def from_flyte_idl(cls, pb2_object): return cls(uri=pb2_object.uri, type=_SchemaType.from_flyte_idl(pb2_object.type)) +class Union(_common.FlyteIdlEntity): + def __init__(self, value, type, tag): + """ + The runtime representation of a tagged union value. See `UnionType` for more details. + + :param flytekit.models.literals.Literal value: + :param flytekit.models.types.UnionType type: + :param int tag: + """ + self._value = value + self._type = type + self._tag = tag + + @property + def value(self): + """ + :rtype: flytekit.models.literals.Literal + """ + return self._value + + @property + def type(self): + """ + :rtype: flytekit.models.types.UnionType + """ + return self._type + + @property + def tag(self): + """ + :rtype: int + """ + return self._tag + + def to_flyte_idl(self): + """ + :rtype: flyteidl.core.literals_pb2.Union + """ + return _literals_pb2.Union(value=self.value.to_flyte_idl(), type=self.type.to_flyte_idl(), tag=self.tag) + + @classmethod + def from_flyte_idl(cls, pb2_object): + """ + :param flyteidl.core.literals_pb2.Schema pb2_object: + :rtype: Schema + """ + return cls(value=Literal.from_flyte_idl(pb2_object.value), type=_UnionType.from_flyte_idl(pb2_object.type), tag=pb2_object.tag) + class LiteralCollection(_common.FlyteIdlEntity): def __init__(self, literals): """ @@ -612,6 +661,7 @@ def __init__( blob: Blob = None, binary: Binary = None, schema: Schema = None, + union: Union = None, none_type: Void = None, error=None, generic: Struct = None, @@ -632,6 +682,7 @@ def __init__( self._blob = blob self._binary = binary self._schema = schema + self._union = union self._none_type = none_type self._error = error self._generic = generic @@ -664,6 +715,13 @@ def schema(self): """ return self._schema + @property + def union(self): + """ + :rtype: Union + """ + return self._union + @property def none_type(self): """ @@ -691,7 +749,7 @@ def value(self): Returns whichever value is set :rtype: T """ - return self.primitive or self.blob or self.binary or self.schema or self.none_type or self.error + return self.primitive or self.blob or self.binary or self.schema or self.union or self.none_type or self.error def to_flyte_idl(self): """ @@ -702,6 +760,7 @@ def to_flyte_idl(self): blob=self.blob.to_flyte_idl() if self.blob is not None else None, binary=self.binary.to_flyte_idl() if self.binary is not None else None, schema=self.schema.to_flyte_idl() if self.schema is not None else None, + union=self.union.to_flyte_idl() if self.union is not None else None none_type=self.none_type.to_flyte_idl() if self.none_type is not None else None, error=self.error if self.error is not None else None, generic=self.generic, @@ -719,6 +778,7 @@ def from_flyte_idl(cls, pb2_object): blob=Blob.from_flyte_idl(pb2_object.blob) if pb2_object.HasField("blob") else None, binary=Binary.from_flyte_idl(pb2_object.binary) if pb2_object.HasField("binary") else None, schema=Schema.from_flyte_idl(pb2_object.schema) if pb2_object.HasField("schema") else None, + union=Union.from_flyte_idl(pb2_object.union) if pb2_object.HasField("union") else None, none_type=Void.from_flyte_idl(pb2_object.none_type) if pb2_object.HasField("none_type") else None, error=pb2_object.error if pb2_object.HasField("error") else None, generic=pb2_object.generic if pb2_object.HasField("generic") else None, diff --git a/flytekit/models/types.py b/flytekit/models/types.py index 7e05fd1593..e013bfd6be 100644 --- a/flytekit/models/types.py +++ b/flytekit/models/types.py @@ -104,21 +104,21 @@ class UnionType(_common.FlyteIdlEntity): Models _types_pb2.UnionType """ - def __init__(self, values: typing.List["LiteralType"]): - self._values = values + def __init__(self, variants: typing.List["LiteralType"]): + self._variants = variants @property - def values(self) -> typing.List["LiteralType"]: - return self._values + def variants(self) -> typing.List["LiteralType"]: + return self._variants def to_flyte_idl(self) -> _types_pb2.UnionType: return _types_pb2.UnionType( - values=[val.to_flyte_idl() if val else None for val in self._values], + variants=[val.to_flyte_idl() if val else None for val in self._variants], ) @classmethod def from_flyte_idl(cls, proto: _types_pb2.UnionType): - return cls(values=[LiteralType.from_flyte_idl(v) for v in proto.values]) + return cls(variants=[LiteralType.from_flyte_idl(v) for v in proto.values]) class LiteralType(_common.FlyteIdlEntity): From a11bfe77f186cf1b91050b7e434a3560d5f2a342 Mon Sep 17 00:00:00 2001 From: Maksim Smolin Date: Thu, 9 Dec 2021 11:43:28 -0800 Subject: [PATCH 008/128] Update to use string tags (part 1) Signed-off-by: maximsmol --- flytekit/core/type_engine.py | 2 +- flytekit/models/literals.py | 19 +++++------------- flytekit/models/types.py | 38 +++++++++++++++++++++++++++++++----- 3 files changed, 39 insertions(+), 20 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 0fa60808b9..b2934714a0 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -646,7 +646,7 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: def guess_python_type(self, literal_type: LiteralType) -> type: if literal_type.union_type is not None: - return typing.Union[tuple(TypeEngine.guess_python_type(v) for v in literal_type.union_type.variants)] + return typing.Union[tuple(TypeEngine.guess_python_type(v.type) for v in literal_type.union_type.variants)] raise ValueError(f"Union transformer cannot reverse {literal_type}") diff --git a/flytekit/models/literals.py b/flytekit/models/literals.py index e45263eca3..e8bb34c441 100644 --- a/flytekit/models/literals.py +++ b/flytekit/models/literals.py @@ -548,16 +548,14 @@ def from_flyte_idl(cls, pb2_object): class Union(_common.FlyteIdlEntity): - def __init__(self, value, type, tag): + def __init__(self, value, tag): """ The runtime representation of a tagged union value. See `UnionType` for more details. :param flytekit.models.literals.Literal value: - :param flytekit.models.types.UnionType type: - :param int tag: + :param str tag: """ self._value = value - self._type = type self._tag = tag @property @@ -567,17 +565,10 @@ def value(self): """ return self._value - @property - def type(self): - """ - :rtype: flytekit.models.types.UnionType - """ - return self._type - @property def tag(self): """ - :rtype: int + :rtype: str """ return self._tag @@ -585,7 +576,7 @@ def to_flyte_idl(self): """ :rtype: flyteidl.core.literals_pb2.Union """ - return _literals_pb2.Union(value=self.value.to_flyte_idl(), type=self.type.to_flyte_idl(), tag=self.tag) + return _literals_pb2.Union(value=self.value.to_flyte_idl(), tag=self.tag) @classmethod def from_flyte_idl(cls, pb2_object): @@ -593,7 +584,7 @@ def from_flyte_idl(cls, pb2_object): :param flyteidl.core.literals_pb2.Schema pb2_object: :rtype: Schema """ - return cls(value=Literal.from_flyte_idl(pb2_object.value), type=_UnionType.from_flyte_idl(pb2_object.type), tag=pb2_object.tag) + return cls(value=Literal.from_flyte_idl(pb2_object.value), tag=pb2_object.tag) class LiteralCollection(_common.FlyteIdlEntity): def __init__(self, literals): diff --git a/flytekit/models/types.py b/flytekit/models/types.py index e013bfd6be..7ba19b10ae 100644 --- a/flytekit/models/types.py +++ b/flytekit/models/types.py @@ -99,26 +99,54 @@ def from_flyte_idl(cls, proto): return cls(columns=[SchemaType.SchemaColumn.from_flyte_idl(c) for c in proto.columns]) +class UnionVariant(_common.FlyteIdlEntity): + """ + Models _types_pb2.UnionVariant + """ + + def __init__(self, type: "LiteralType", tag: str): + self._type = type + self._tag = tag + + @property + def type(self) -> "LiteralType": + return self._type + + @property + def tag(self) -> str: + return self._tag + + def to_flyte_idl(self) -> _types_pb2.UnionVariant: + return _types_pb2.UnionVariant( + type=self.type, + tag=self.tag + ) + + @classmethod + def from_flyte_idl(cls, proto: _types_pb2.UnionVariant): + return cls(type=LiteralType.from_flyte_idl(proto.type), tag=proto.tag) + + class UnionType(_common.FlyteIdlEntity): """ Models _types_pb2.UnionType """ - def __init__(self, variants: typing.List["LiteralType"]): + def __init__(self, variants: typing.List["UnionVariant"]): self._variants = variants @property - def variants(self) -> typing.List["LiteralType"]: + def variants(self) -> typing.List["UnionVariant"]: return self._variants - def to_flyte_idl(self) -> _types_pb2.UnionType: - return _types_pb2.UnionType( + def to_flyte_idl(self) -> _types_pb2.UnionVariant: + return _types_pb2.UnionVariant( variants=[val.to_flyte_idl() if val else None for val in self._variants], ) @classmethod def from_flyte_idl(cls, proto: _types_pb2.UnionType): - return cls(variants=[LiteralType.from_flyte_idl(v) for v in proto.values]) + return cls(variants=[UnionVariant.from_flyte_idl(v) for v in proto.variants]) class LiteralType(_common.FlyteIdlEntity): From 38b9dfa1946058aa53efd27f46f1ad087ff4f774 Mon Sep 17 00:00:00 2001 From: Maksim Smolin Date: Thu, 9 Dec 2021 16:06:13 -0800 Subject: [PATCH 009/128] Working implementation, update tests Signed-off-by: maximsmol --- .github/workflows/pythonbuild.yml | 125 +++++++++---------- flytekit/core/type_engine.py | 92 +++++++++++--- flytekit/models/literals.py | 2 +- tests/flytekit/unit/core/test_type_engine.py | 27 ++-- 4 files changed, 156 insertions(+), 90 deletions(-) diff --git a/.github/workflows/pythonbuild.yml b/.github/workflows/pythonbuild.yml index b83bd79bfd..ba633c0a19 100644 --- a/.github/workflows/pythonbuild.yml +++ b/.github/workflows/pythonbuild.yml @@ -3,7 +3,7 @@ name: Build on: push: branches: - - master + - master pull_request: jobs: @@ -15,42 +15,42 @@ jobs: python-version: [3.7, 3.8, 3.9] spark-version-suffix: ["", "-spark2"] exclude: - - python-version: 3.8 - spark-version-suffix: "-spark2" - - python-version: 3.9 - spark-version-suffix: "-spark2" + - python-version: 3.8 + spark-version-suffix: "-spark2" + - python-version: 3.9 + spark-version-suffix: "-spark2" steps: - - uses: actions/checkout@v2 - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v2 - with: - python-version: ${{ matrix.python-version }} - - name: Cache pip - uses: actions/cache@v2 - with: - # This path is specific to Ubuntu - path: ~/.cache/pip - # Look to see if there is a cache hit for the corresponding requirements files - key: ${{ format('{0}-pip-{1}', runner.os, hashFiles('dev-requirements.txt', format('requirements{0}.txt', matrix.spark-version-suffix))) }} - - name: Install dependencies - run: | - python -m pip install --upgrade pip==21.2.4 setuptools wheel - make setup${{ matrix.spark-version-suffix }} - git clone https://github.com/flyteorg/flyteidl.git && cd flyteidl && git checkout union_type && pip install . && cd .. - pip freeze - - name: Test with coverage - run: | - coverage run -m pytest tests/flytekit/unit - - name: Integration Tests with coverage - # https://github.com/actions/runner/issues/241#issuecomment-577360161 - shell: 'script -q -e -c "bash {0}"' - run: | - python -m pip install awscli - coverage run --append -m pytest tests/flytekit/integration - - name: Codecov - uses: codecov/codecov-action@v1 - with: - fail_ci_if_error: true # optional (default = false) + - uses: actions/checkout@v2 + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v2 + with: + python-version: ${{ matrix.python-version }} + - name: Cache pip + uses: actions/cache@v2 + with: + # This path is specific to Ubuntu + path: ~/.cache/pip + # Look to see if there is a cache hit for the corresponding requirements files + key: ${{ format('{0}-pip-{1}', runner.os, hashFiles('dev-requirements.txt', format('requirements{0}.txt', matrix.spark-version-suffix))) }} + - name: Install dependencies + run: | + python -m pip install --upgrade pip==21.2.4 setuptools wheel + make setup${{ matrix.spark-version-suffix }} + git clone https://github.com/maximsmol/flyteidl.git && cd flyteidl && git checkout maximsmol-union-types && pip install . && cd .. + pip freeze + - name: Test with coverage + run: | + coverage run -m pytest tests/flytekit/unit + - name: Integration Tests with coverage + # https://github.com/actions/runner/issues/241#issuecomment-577360161 + shell: 'script -q -e -c "bash {0}"' + run: | + python -m pip install awscli + coverage run --append -m pytest tests/flytekit/integration + - name: Codecov + uses: codecov/codecov-action@v1 + with: + fail_ci_if_error: true # optional (default = false) build-plugins: runs-on: ubuntu-latest @@ -75,31 +75,31 @@ jobs: - flytekit-snowflake - flytekit-modin steps: - - uses: actions/checkout@v2 - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v2 - with: - python-version: ${{ matrix.python-version }} - - name: Cache pip - uses: actions/cache@v2 - with: - # This path is specific to Ubuntu - path: ~/.cache/pip - # Look to see if there is a cache hit for the corresponding requirements files - key: ${{ format('{0}-pip-{1}', runner.os, hashFiles('dev-requirements.txt', format('plugins/{0}/requirements.txt', matrix.plugin-names ))) }} - - name: Install dependencies - run: | - python -m pip install --upgrade pip==21.2.4 setuptools wheel - make setup - cd plugins/${{ matrix.plugin-names }} - pip install -e . - pip install --no-deps -U https://github.com/flyteorg/flytekit/archive/${{ github.sha }}.zip#egg=flytekit - git clone https://github.com/flyteorg/flyteidl.git && cd flyteidl && git checkout union_type && pip install . && cd .. - pip freeze - - name: Test with coverage - run: | - cd plugins/${{ matrix.plugin-names }} - coverage run -m pytest tests + - uses: actions/checkout@v2 + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v2 + with: + python-version: ${{ matrix.python-version }} + - name: Cache pip + uses: actions/cache@v2 + with: + # This path is specific to Ubuntu + path: ~/.cache/pip + # Look to see if there is a cache hit for the corresponding requirements files + key: ${{ format('{0}-pip-{1}', runner.os, hashFiles('dev-requirements.txt', format('plugins/{0}/requirements.txt', matrix.plugin-names ))) }} + - name: Install dependencies + run: | + python -m pip install --upgrade pip==21.2.4 setuptools wheel + make setup + cd plugins/${{ matrix.plugin-names }} + pip install -e . + pip install --no-deps -U https://github.com/flyteorg/flytekit/archive/${{ github.sha }}.zip#egg=flytekit + git clone https://github.com/flyteorg/flyteidl.git && cd flyteidl && git checkout union_type && pip install . && cd .. + pip freeze + - name: Test with coverage + run: | + cd plugins/${{ matrix.plugin-names }} + coverage run -m pytest tests lint: runs-on: ubuntu-latest @@ -126,8 +126,7 @@ jobs: - name: ShellCheck uses: ludeeus/action-shellcheck@master with: - ignore: - boilerplate + ignore: boilerplate docs: runs-on: ubuntu-latest diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index b2934714a0..b21d8c4aec 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -8,6 +8,7 @@ import mimetypes import typing from abc import ABC, abstractmethod +from re import L from typing import NamedTuple, Optional, Type, cast from dataclasses_json import DataClassJsonMixin, dataclass_json @@ -27,12 +28,14 @@ from flytekit.models import interface as _interface_models from flytekit.models import types as _type_models from flytekit.models.core import types as _core_types -from flytekit.models.literals import Literal, LiteralCollection, LiteralMap, Primitive, Scalar, Union -from flytekit.models.types import LiteralType, SimpleType, UnionType +from flytekit.models.literals import Literal, LiteralCollection, LiteralMap, Primitive, Scalar, Union, Void +from flytekit.models.types import LiteralType, SimpleType, UnionType, UnionVariant T = typing.TypeVar("T") DEFINITIONS = "definitions" +class TypeTransformerFailedError(RuntimeError): + ... class TypeTransformer(typing.Generic[T]): """ @@ -126,6 +129,7 @@ def __init__( from_literal_transformer: typing.Callable[[Literal], T], ): super().__init__(name, t) + self._type = t self._lt = lt self._to_literal_transformer = to_literal_transformer self._from_literal_transformer = from_literal_transformer @@ -134,6 +138,8 @@ def get_literal_type(self, t: Type[T] = None) -> LiteralType: return self._lt def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], expected: LiteralType) -> Literal: + if type(python_val) != self._type: + raise TypeTransformerFailedError(f"Expected value of type {self._type} but got type {type(python_val)}") return self._to_literal_transformer(python_val) def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[T]) -> T: @@ -234,12 +240,12 @@ def get_literal_type(self, t: Type[T]) -> LiteralType: def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], expected: LiteralType) -> Literal: if not dataclasses.is_dataclass(python_val): - raise AssertionError( + raise TypeTransformerFailedError( f"{type(python_val)} is not of type @dataclass, only Dataclasses are supported for " f"user defined datatypes in Flytekit" ) if not issubclass(type(python_val), DataClassJsonMixin): - raise AssertionError( + raise TypeTransformerFailedError( f"Dataclass {python_type} should be decorated with @dataclass_json to be " f"serialized correctly" ) return Literal( @@ -315,7 +321,10 @@ def get_literal_type(self, t: Type[T]) -> LiteralType: def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], expected: LiteralType) -> Literal: struct = Struct() - struct.update(_MessageToDict(python_val)) + try: + struct.update(_MessageToDict(python_val)) + except: + raise TypeTransformerFailedError("Failed to convert to generic protobuf struct") return Literal(scalar=Scalar(generic=struct)) def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[T]) -> T: @@ -581,6 +590,9 @@ def get_literal_type(self, t: Type[T]) -> Optional[LiteralType]: raise ValueError(f"Type of Generic List type is not supported, {e}") def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], expected: LiteralType) -> Literal: + if type(python_val) != list: + raise TypeTransformerFailedError("Expected a list") + t = self.get_sub_type(python_type) lit_list = [TypeEngine.to_literal(ctx, x, t, expected.collection_type) for x in python_val] # type: ignore return Literal(collection=LiteralCollection(literals=lit_list)) @@ -606,22 +618,34 @@ def __init__(self): def get_literal_type(self, t: Type[T]) -> Optional[LiteralType]: try: - sub_type = [TypeEngine.to_literal_type(v) for v in typing.get_args(t)] - return _type_models.LiteralType(union_type=UnionType(sub_type)) + trans = [(TypeEngine.get_transformer(x), x) for x in typing.get_args(t)] + variants = [UnionVariant(t.get_literal_type(x), t.name) for (t, x) in trans] + return _type_models.LiteralType(union_type=UnionType(variants)) except Exception as e: raise ValueError(f"Type of Generic Union type is not supported, {e}") def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], expected: LiteralType) -> Literal: - for idx, t in enumerate(typing.get_args(python_type)): + found_res = False + res = None + res_tag = None + for t in typing.get_args(python_type): try: - val = TypeEngine.to_literal(ctx, python_val, t, expected) - except Exception as e: + trans = TypeEngine.get_transformer(t) + + res = trans.to_literal(ctx, python_val, t, expected) + res_tag = trans.name + if found_res: + # Should really never happen, sanity check + raise AssertionError(f"Ambiguous choice of variant for union type") + found_res = True + except TypeTransformerFailedError as e: logger.debug(f"Failed to convert from {python_val} to {t}", e) continue - return Literal(scalar=Scalar(union=Union(value=val, type=expected, tag=idx))) + if found_res: + return Literal(scalar=Scalar(union=Union(value=res, tag=res_tag))) - raise TypeError(f"Cannot convert from {python_val} to {python_type}") + raise TypeTransformerFailedError(f"Cannot convert from {python_val} to {python_type}") def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[T]) -> Optional[typing.Any]: for x in typing.get_args(expected_python_type): @@ -634,15 +658,37 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: break + union_tag = None + if lv.scalar is not None and lv.scalar.union is not None: + union_tag = lv.scalar.union.tag + + found_res = False + res = None for v in typing.get_args(expected_python_type): try: - val = TypeEngine.to_python_value(ctx, lv, v) - if val: - return val + trans = TypeEngine.get_transformer(v) + if union_tag: + if trans.name != union_tag: + continue + + assert lv.scalar is not None # type checker + assert lv.scalar.union is not None # type checker + + res = trans.to_python_value(ctx, lv.scalar.union.value, v) + found_res = True + break + else: + res = trans.to_python_value(ctx, lv, v) + if found_res: + raise TypeError(f"Ambiguous choice of vaariant for union type") + found_res = True except Exception as e: # todo(maximsmol): separate programmer errors from type conversion failures logger.debug(f"Failed to convert from {lv} to {v}", e) - raise TypeError(f"Cannot convert from {lv} to {expected_python_type}") + if found_res: + return res + + raise TypeError(f"Cannot convert from {lv} to {expected_python_type} (using tag {union_tag})") def guess_python_type(self, literal_type: LiteralType) -> type: if literal_type.union_type is not None: @@ -694,6 +740,9 @@ def get_literal_type(self, t: Type[dict]) -> LiteralType: def to_literal( self, ctx: FlyteContext, python_val: typing.Any, python_type: Type[dict], expected: LiteralType ) -> Literal: + if type(python_val) != dict: + raise TypeTransformerFailedError("Expected a dict") + if expected and expected.simple and expected.simple == SimpleType.STRUCT: return self.dict_to_generic_literal(python_val) @@ -821,6 +870,11 @@ def get_literal_type(self, t: Type[T]) -> LiteralType: return LiteralType(enum_type=_core_types.EnumType(values=values)) def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], expected: LiteralType) -> Literal: + if type(python_val).__class__ != enum.EnumMeta: + raise TypeTransformerFailedError("Expected an enum") + if type(python_val.value) != str: + raise TypeTransformerFailedError("Only string-valued enums are sujpportedd") + return Literal(scalar=Scalar(primitive=Primitive(string_value=python_val.value))) # type: ignore def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[T]) -> T: @@ -958,12 +1012,12 @@ def _register_default_type_transformers(): TypeEngine.register( SimpleTransformer( "none", - None, + type(None), _type_models.LiteralType(simple=_type_models.SimpleType.NONE), - lambda x: None, + lambda x: Literal(scalar=Scalar(none_type=Void())), lambda x: None, ), - [type(None)], + [None], ) TypeEngine.register(ListTransformer()) TypeEngine.register(UnionTransformer()) diff --git a/flytekit/models/literals.py b/flytekit/models/literals.py index e8bb34c441..166b03e6e6 100644 --- a/flytekit/models/literals.py +++ b/flytekit/models/literals.py @@ -751,7 +751,7 @@ def to_flyte_idl(self): blob=self.blob.to_flyte_idl() if self.blob is not None else None, binary=self.binary.to_flyte_idl() if self.binary is not None else None, schema=self.schema.to_flyte_idl() if self.schema is not None else None, - union=self.union.to_flyte_idl() if self.union is not None else None + union=self.union.to_flyte_idl() if self.union is not None else None, none_type=self.none_type.to_flyte_idl() if self.none_type is not None else None, error=self.error if self.error is not None else None, generic=self.generic, diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 1f77828373..526b9726de 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -25,7 +25,7 @@ ) from flytekit.models import types as model_types from flytekit.models.core.types import BlobType -from flytekit.models.literals import Blob, BlobMetadata, Literal, LiteralCollection, LiteralMap, Primitive, Scalar +from flytekit.models.literals import Blob, BlobMetadata, Literal, LiteralCollection, LiteralMap, Primitive, Scalar, Void from flytekit.models.types import LiteralType, SimpleType from flytekit.types.directory.types import FlyteDirectory from flytekit.types.file import JPEGImageFile @@ -349,7 +349,7 @@ def test_guessing_basic(): lt = model_types.LiteralType(simple=model_types.SimpleType.NONE) pt = TypeEngine.guess_python_type(lt) - assert pt is None + assert pt is type(None) def test_guessing_containers(): @@ -547,36 +547,49 @@ def test_enum_type(): TypeEngine.to_literal_type(UnsupportedEnumValues) +def union_type_tags_unique(t: LiteralType): + seen = set() + for x in t.union_type.variants: + if x.tag in seen: + return False + seen.add(x.tag) + + return True + + def test_union_type(): pt = typing.Union[str, int] lt = TypeEngine.to_literal_type(pt) - assert lt.union_type.values == [LiteralType(simple=SimpleType.STRING), LiteralType(simple=SimpleType.INTEGER)] + assert [x.type for x in lt.union_type.variants] == [LiteralType(simple=SimpleType.STRING), LiteralType(simple=SimpleType.INTEGER)] + assert union_type_tags_unique(lt) ctx = FlyteContextManager.current_context() lv = TypeEngine.to_literal(ctx, 3, pt, lt) v = TypeEngine.to_python_value(ctx, lv, pt) - assert lv.scalar.primitive.integer == 3 + assert lv.scalar.union.value.scalar.primitive.integer == 3 assert v == 3 lv = TypeEngine.to_literal(ctx, "hello", pt, lt) v = TypeEngine.to_python_value(ctx, lv, pt) - assert lv.scalar.primitive.string_value == "hello" + assert lv.scalar.union.value.scalar.primitive.string_value == "hello" assert v == "hello" def test_optional_type(): pt = typing.Optional[int] lt = TypeEngine.to_literal_type(pt) - assert lt.union_type.values == [LiteralType(simple=SimpleType.INTEGER), LiteralType(simple=SimpleType.NONE)] + assert [x.type for x in lt.union_type.variants] == [LiteralType(simple=SimpleType.INTEGER), LiteralType(simple=SimpleType.NONE)] + assert union_type_tags_unique(lt) ctx = FlyteContextManager.current_context() lv = TypeEngine.to_literal(ctx, 3, pt, lt) v = TypeEngine.to_python_value(ctx, lv, pt) - assert lv.scalar.primitive.integer == 3 + assert lv.scalar.union.value.scalar.primitive.integer == 3 assert v == 3 lv = TypeEngine.to_literal(ctx, None, pt, lt) v = TypeEngine.to_python_value(ctx, lv, pt) + assert lv.scalar.union.value.scalar.none_type == Void() assert v is None From 2ca507b8485205d1b01391dabc62ab8809b51edb Mon Sep 17 00:00:00 2001 From: Maksim Smolin Date: Thu, 9 Dec 2021 16:50:43 -0800 Subject: [PATCH 010/128] Fixes, more tests Signed-off-by: maximsmol --- flytekit/core/type_engine.py | 36 ++++--- tests/flytekit/unit/core/test_type_engine.py | 108 +++++++++++++++++++ 2 files changed, 129 insertions(+), 15 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index b21d8c4aec..d585544aaf 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -143,7 +143,14 @@ def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], exp return self._to_literal_transformer(python_val) def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[T]) -> T: - return self._from_literal_transformer(lv) + if expected_python_type != self._type: + raise TypeTransformerFailedError(f"Cannot convert to type {expected_python_type}, only {self._type} is supported") + + try: # todo(maximsmol): this is quite ugly and each transformer should really check their Literal + return self._from_literal_transformer(lv) + except AttributeError: + # Assume that this is because a property on `lv` was None + raise TypeTransformerFailedError(f"Cannot convert literal {lv}") def guess_python_type(self, literal_type: LiteralType) -> Type[T]: if literal_type.simple is not None and literal_type.simple == self._lt.simple: @@ -648,22 +655,13 @@ def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], exp raise TypeTransformerFailedError(f"Cannot convert from {python_val} to {python_type}") def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[T]) -> Optional[typing.Any]: - for x in typing.get_args(expected_python_type): - if x is type(None): - if lv is None: # todo(maximsmol): why would this ever be None and not Void? - return None - - if lv.scalar is not None and lv.scalar.none_type is not None: - return None - - break - union_tag = None if lv.scalar is not None and lv.scalar.union is not None: union_tag = lv.scalar.union.tag found_res = False res = None + res_tag = None for v in typing.get_args(expected_python_type): try: trans = TypeEngine.get_transformer(v) @@ -675,14 +673,18 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: assert lv.scalar.union is not None # type checker res = trans.to_python_value(ctx, lv.scalar.union.value, v) + res_tag = trans.name found_res = True break else: res = trans.to_python_value(ctx, lv, v) if found_res: - raise TypeError(f"Ambiguous choice of vaariant for union type") + raise TypeError( + "Ambiguous choice of variant for union type. " + + f"Both {res_tag} and {trans.name} transformers match") + res_tag = trans.name found_res = True - except Exception as e: # todo(maximsmol): separate programmer errors from type conversion failures + except TypeTransformerFailedError as e: # todo(maximsmol): separate programmer errors from type conversion failures logger.debug(f"Failed to convert from {lv} to {v}", e) if found_res: @@ -945,8 +947,12 @@ def _check_and_covert_float(lv: Literal) -> float: return lv.scalar.primitive.float_value elif lv.scalar.primitive.integer is not None: return float(lv.scalar.primitive.integer) - raise RuntimeError(f"Cannot convert literal {lv} to float") + raise TypeTransformerFailedError(f"Cannot convert literal {lv} to float") +def _check_and_convert_void(lv: Literal) -> None: + if lv.scalar.none_type is None: + raise TypeTransformerFailedError(f"Cannot conver literal {lv} to None") + return None def _register_default_type_transformers(): TypeEngine.register( @@ -1015,7 +1021,7 @@ def _register_default_type_transformers(): type(None), _type_models.LiteralType(simple=_type_models.SimpleType.NONE), lambda x: Literal(scalar=Scalar(none_type=Void())), - lambda x: None, + lambda x: _check_and_convert_void(x) ), [None], ) diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 526b9726de..5b57ab298d 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -13,6 +13,7 @@ from marshmallow_jsonschema import JSONSchema from flytekit.common.exceptions import user as user_exceptions +from flytekit.common.types import primitives from flytekit.core.context_manager import FlyteContext, FlyteContextManager from flytekit.core.type_engine import ( DataclassTransformer, @@ -20,6 +21,8 @@ ListTransformer, SimpleTransformer, TypeEngine, + TypeTransformer, + TypeTransformerFailedError, convert_json_schema_to_python_class, dataclass_from_dict, ) @@ -33,6 +36,7 @@ from flytekit.types.pickle import FlytePickle from flytekit.types.pickle.pickle import FlytePickleTransformer +T = typing.TypeVar("T") def test_type_engine(): t = int @@ -592,6 +596,110 @@ def test_optional_type(): assert lv.scalar.union.value.scalar.none_type == Void() assert v is None +def test_union_from_unambiguous_literal(): + pt = typing.Optional[int] + lt = TypeEngine.to_literal_type(pt) + assert [x.type for x in lt.union_type.variants] == [LiteralType(simple=SimpleType.INTEGER), LiteralType(simple=SimpleType.NONE)] + assert union_type_tags_unique(lt) + + ctx = FlyteContextManager.current_context() + lv = TypeEngine.to_literal(ctx, 3, int, primitives.Integer.to_flyte_literal_type()) + assert lv.scalar.primitive.integer == 3 + + v = TypeEngine.to_python_value(ctx, lv, pt) + assert v == 3 + +def test_union_custom_transformer(): + class MyInt: + def __init__(self, x: int): + self.val = x + + def __eq__(self, other): + if not isinstance(other, MyInt): + return False + return other.val == self.val + + TypeEngine.register( + SimpleTransformer( + "MyInt", + MyInt, + primitives.Integer.to_flyte_literal_type(), + lambda x: Literal(scalar=Scalar(primitive=Primitive(integer=x.val))), + lambda x: MyInt(x.scalar.primitive.integer), + ) + ) + + pt = typing.Union[int, MyInt] + lt = TypeEngine.to_literal_type(pt) + assert [x.type for x in lt.union_type.variants] == [LiteralType(simple=SimpleType.INTEGER), LiteralType(simple=SimpleType.INTEGER)] + assert union_type_tags_unique(lt) + + ctx = FlyteContextManager.current_context() + lv = TypeEngine.to_literal(ctx, 3, pt, lt) + v = TypeEngine.to_python_value(ctx, lv, pt) + assert lv.scalar.union.value.scalar.primitive.integer == 3 + assert v == 3 + + lv = TypeEngine.to_literal(ctx, MyInt(10), pt, lt) + v = TypeEngine.to_python_value(ctx, lv, pt) + assert lv.scalar.union.value.scalar.primitive.integer == 10 + assert v == MyInt(10) + + lv = TypeEngine.to_literal(ctx, 4, int, primitives.Integer.to_flyte_literal_type()) + assert lv.scalar.primitive.integer == 4 + try: + TypeEngine.to_python_value(ctx, lv, pt) + except TypeError as e: + assert "Ambiguous choice of variant" in str(e) + + del TypeEngine._REGISTRY[MyInt] + +def test_union_custom_transformer_sanity_check(): + class UnsignedInt: + def __init__(self, x: int): + self.val = x + + def __eq__(self, other): + if not isinstance(other, UnsignedInt): + return False + return other.val == self.val + + class UnsignedIntTransformer(TypeTransformer[UnsignedInt]): + def __init__(self): + super().__init__("UnsignedInt", UnsignedInt) + + def get_literal_type(self, t: typing.Type[T]) -> LiteralType: + return primitives.Integer.to_flyte_literal_type() + + def to_literal(self, ctx: FlyteContext, python_val: T, python_type: typing.Type[T], expected: LiteralType) -> Literal: + if type(python_val) != int: + raise TypeTransformerFailedError("Expected an integer") + + if python_val < 0: + raise TypeTransformerFailedError("Expected a non-negative integer") + + return Literal(scalar=Scalar(primitive=Primitive(integer=python_val))) + + def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: typing.Type[T]) -> T: + val = lv.scalar.primitive.integer + return UnsignedInt(0 if val < 0 else val) + + TypeEngine.register(UnsignedIntTransformer()) + + + pt = typing.Union[int, UnsignedInt] + lt = TypeEngine.to_literal_type(pt) + assert [x.type for x in lt.union_type.variants] == [LiteralType(simple=SimpleType.INTEGER), LiteralType(simple=SimpleType.INTEGER)] + assert union_type_tags_unique(lt) + + ctx = FlyteContextManager.current_context() + try: + TypeEngine.to_literal(ctx, 3, pt, lt) + except AssertionError as e: + assert str(e) == "Ambiguous choice of variant for union type" + + del TypeEngine._REGISTRY[UnsignedInt] + @pytest.mark.parametrize( "python_value,python_types,expected_literal_map", From deb53a5dbfc14d8599925993cb7a8f5c5429b49f Mon Sep 17 00:00:00 2001 From: Maksim Smolin Date: Fri, 10 Dec 2021 11:09:41 -0800 Subject: [PATCH 011/128] Fix incorrect unwrapped literal-union matching, update test Signed-off-by: maximsmol --- flytekit/core/type_engine.py | 7 +++++-- tests/flytekit/unit/core/test_type_engine.py | 4 ++-- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index d585544aaf..a30fed98da 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -147,8 +147,11 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: raise TypeTransformerFailedError(f"Cannot convert to type {expected_python_type}, only {self._type} is supported") try: # todo(maximsmol): this is quite ugly and each transformer should really check their Literal - return self._from_literal_transformer(lv) - except AttributeError: + res = self._from_literal_transformer(lv) + if type(res) != self._type: + raise TypeTransformerFailedError(f"Cannot convert literal {lv} to {self._type}") + return res + except AttributeError as e: # Assume that this is because a property on `lv` was None raise TypeTransformerFailedError(f"Cannot convert literal {lv}") diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 5b57ab298d..c5a523dac0 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -597,9 +597,9 @@ def test_optional_type(): assert v is None def test_union_from_unambiguous_literal(): - pt = typing.Optional[int] + pt = typing.Union[str, int] lt = TypeEngine.to_literal_type(pt) - assert [x.type for x in lt.union_type.variants] == [LiteralType(simple=SimpleType.INTEGER), LiteralType(simple=SimpleType.NONE)] + assert [x.type for x in lt.union_type.variants] == [LiteralType(simple=SimpleType.STRING), LiteralType(simple=SimpleType.INTEGER)] assert union_type_tags_unique(lt) ctx = FlyteContextManager.current_context() From 154c8a0e7eb1e7fb6e6d5f363056661a4cd95649 Mon Sep 17 00:00:00 2001 From: Maksim Smolin Date: Fri, 10 Dec 2021 11:38:52 -0800 Subject: [PATCH 012/128] Fix duplicate tag handling, add tests for collections containing unions Signed-off-by: maximsmol --- flytekit/core/type_engine.py | 7 ++- tests/flytekit/unit/core/test_type_engine.py | 48 +++++++++++++++++++- 2 files changed, 52 insertions(+), 3 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index a30fed98da..7eaf2a4b1e 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -67,7 +67,7 @@ def type_assertions_enabled(self) -> bool: def assert_type(self, t: Type[T], v: T): if not hasattr(t, "__origin__") and not isinstance(v, t): - raise TypeError(f"Type of Val '{v}' is not an instance of {t}") + raise TypeTransformerFailedError(f"Type of Val '{v}' is not an instance of {t}") @abstractmethod def get_literal_type(self, t: Type[T]) -> LiteralType: @@ -677,8 +677,11 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: res = trans.to_python_value(ctx, lv.scalar.union.value, v) res_tag = trans.name + if found_res: + raise TypeError( + "Ambiguous choice of variant for union type. " + + f"Both {res_tag} and {trans.name} transformers match") found_res = True - break else: res = trans.to_python_value(ctx, lv, v) if found_res: diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index c5a523dac0..5dc053ad6c 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -570,11 +570,13 @@ def test_union_type(): ctx = FlyteContextManager.current_context() lv = TypeEngine.to_literal(ctx, 3, pt, lt) v = TypeEngine.to_python_value(ctx, lv, pt) + assert lv.scalar.union.tag == "int" assert lv.scalar.union.value.scalar.primitive.integer == 3 assert v == 3 lv = TypeEngine.to_literal(ctx, "hello", pt, lt) v = TypeEngine.to_python_value(ctx, lv, pt) + assert lv.scalar.union.tag == "str" assert lv.scalar.union.value.scalar.primitive.string_value == "hello" assert v == "hello" @@ -588,11 +590,13 @@ def test_optional_type(): ctx = FlyteContextManager.current_context() lv = TypeEngine.to_literal(ctx, 3, pt, lt) v = TypeEngine.to_python_value(ctx, lv, pt) + assert lv.scalar.union.tag == "int" assert lv.scalar.union.value.scalar.primitive.integer == 3 assert v == 3 lv = TypeEngine.to_literal(ctx, None, pt, lt) v = TypeEngine.to_python_value(ctx, lv, pt) + assert lv.scalar.union.tag == "none" assert lv.scalar.union.value.scalar.none_type == Void() assert v is None @@ -637,11 +641,13 @@ def __eq__(self, other): ctx = FlyteContextManager.current_context() lv = TypeEngine.to_literal(ctx, 3, pt, lt) v = TypeEngine.to_python_value(ctx, lv, pt) + assert lv.scalar.union.tag == "int" assert lv.scalar.union.value.scalar.primitive.integer == 3 assert v == 3 lv = TypeEngine.to_literal(ctx, MyInt(10), pt, lt) v = TypeEngine.to_python_value(ctx, lv, pt) + assert lv.scalar.union.tag == "MyInt" assert lv.scalar.union.value.scalar.primitive.integer == 10 assert v == MyInt(10) @@ -664,6 +670,7 @@ def __eq__(self, other): return False return other.val == self.val + # This transformer will not work in the implicit wrapping case class UnsignedIntTransformer(TypeTransformer[UnsignedInt]): def __init__(self): super().__init__("UnsignedInt", UnsignedInt) @@ -701,6 +708,45 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: del TypeEngine._REGISTRY[UnsignedInt] +def test_union_of_lists(): + pt = typing.Union[typing.List[int], typing.List[str]] + lt = TypeEngine.to_literal_type(pt) + assert [x.type for x in lt.union_type.variants] == [ + LiteralType(collection_type=LiteralType(simple=SimpleType.INTEGER)), + LiteralType(collection_type=LiteralType(simple=SimpleType.STRING)), + ] + assert not union_type_tags_unique(lt) # tags are deliberately NOT unique + + ctx = FlyteContextManager.current_context() + lv = TypeEngine.to_literal(ctx, ["hello", "world"], pt, lt) + v = TypeEngine.to_python_value(ctx, lv, pt) + assert lv.scalar.union.tag == "Typed List" + assert [x.scalar.primitive.string_value for x in lv.scalar.union.value.collection.literals] == ["hello", "world"] + assert v == ["hello", "world"] + + lv = TypeEngine.to_literal(ctx, [1, 3], pt, lt) + v = TypeEngine.to_python_value(ctx, lv, pt) + assert lv.scalar.union.tag == "Typed List" + assert [x.scalar.primitive.integer for x in lv.scalar.union.value.collection.literals] == [1, 3] + assert v == [1, 3] + + +def test_list_of_unions(): + pt = typing.List[typing.Union[str, int]] + lt = TypeEngine.to_literal_type(pt) + # todo(maximsmol): seems like the order here is non-deterministic + assert [x.type for x in lt.collection_type.union_type.variants] == [ + LiteralType(simple=SimpleType.STRING), + LiteralType(simple=SimpleType.INTEGER), + ] + assert union_type_tags_unique(lt.collection_type) # tags are deliberately NOT unique + + ctx = FlyteContextManager.current_context() + lv = TypeEngine.to_literal(ctx, ["hello", 123, "world"], pt, lt) + v = TypeEngine.to_python_value(ctx, lv, pt) + assert [x.scalar.union.tag for x in lv.collection.literals] == ["str", "int", "str"] + assert v == ["hello", 123, "world"] + @pytest.mark.parametrize( "python_value,python_types,expected_literal_map", [ @@ -786,5 +832,5 @@ def test_dict_to_literal_map_with_wrong_input_type(): ctx = FlyteContext.current_context() input = {"a": 1} guessed_python_types = {"a": str} - with pytest.raises(user_exceptions.FlyteTypeException): + with pytest.raises(TypeTransformerFailedError): TypeEngine.dict_to_literal_map(ctx, input, guessed_python_types) From b5772156a9ed20d7d1f3978a339218dc0bfd9b1d Mon Sep 17 00:00:00 2001 From: Maksim Smolin Date: Fri, 10 Dec 2021 11:54:23 -0800 Subject: [PATCH 013/128] Fix type hint test Signed-off-by: maximsmol --- flytekit/core/type_engine.py | 14 +++++++------- flytekit/types/file/file.py | 8 ++++---- flytekit/types/schema/types.py | 9 ++++++--- tests/flytekit/unit/core/test_type_hints.py | 17 ++++++++++++++++- 4 files changed, 33 insertions(+), 15 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 7eaf2a4b1e..0720ab1b34 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -34,7 +34,7 @@ T = typing.TypeVar("T") DEFINITIONS = "definitions" -class TypeTransformerFailedError(RuntimeError): +class TypeTransformerFailedError(TypeError): ... class TypeTransformer(typing.Generic[T]): @@ -295,12 +295,12 @@ def _fix_dataclass_int(self, dc_type: Type[DataClassJsonMixin], dc: DataClassJso def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[T]) -> T: if not dataclasses.is_dataclass(expected_python_type): - raise AssertionError( + raise TypeTransformerFailedError( f"{expected_python_type} is not of type @dataclass, only Dataclasses are supported for " f"user defined datatypes in Flytekit" ) if not issubclass(expected_python_type, DataClassJsonMixin): - raise AssertionError( + raise TypeTransformerFailedError( f"Dataclass {expected_python_type} should be decorated with @dataclass_json to be " f"serialized correctly" ) @@ -339,7 +339,7 @@ def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], exp def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[T]) -> T: if not (lv and lv.scalar and lv.scalar.generic is not None): - raise AssertionError("Can only convert a generic literal to a Protobuf") + raise TypeTransformerFailedError("Can only convert a generic literal to a Protobuf") pb_obj = expected_python_type() dictionary = _MessageToDict(lv.scalar.generic) @@ -464,7 +464,7 @@ def to_literal(cls, ctx: FlyteContext, python_val: typing.Any, python_type: Type Converts a python value of a given type and expected ``LiteralType`` into a resolved ``Literal`` value. """ if python_val is None and expected.union_type is None: - raise AssertionError(f"Python value cannot be None, expected {python_type}/{expected}") + raise TypeTransformerFailedError(f"Python value cannot be None, expected {python_type}/{expected}") transformer = cls.get_transformer(python_type) if transformer.type_assertions_enabled: transformer.assert_type(python_type, python_val) @@ -646,7 +646,7 @@ def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], exp res_tag = trans.name if found_res: # Should really never happen, sanity check - raise AssertionError(f"Ambiguous choice of variant for union type") + raise TypeError(f"Ambiguous choice of variant for union type") found_res = True except TypeTransformerFailedError as e: logger.debug(f"Failed to convert from {python_val} to {t}", e) @@ -874,7 +874,7 @@ def __init__(self): def get_literal_type(self, t: Type[T]) -> LiteralType: values = [v.value for v in t] # type: ignore if not isinstance(values[0], str): - raise AssertionError("Only EnumTypes with value of string are supported") + raise TypeTransformerFailedError("Only EnumTypes with value of string are supported") return LiteralType(enum_type=_core_types.EnumType(values=values)) def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], expected: LiteralType) -> Literal: diff --git a/flytekit/types/file/file.py b/flytekit/types/file/file.py index 1fbcee049a..2421823d6c 100644 --- a/flytekit/types/file/file.py +++ b/flytekit/types/file/file.py @@ -5,7 +5,7 @@ import typing from flytekit.core.context_manager import FlyteContext -from flytekit.core.type_engine import TypeEngine, TypeTransformer +from flytekit.core.type_engine import TypeEngine, TypeTransformer, TypeTransformerFailedError from flytekit.loggers import logger from flytekit.models.core.types import BlobType from flytekit.models.literals import Blob, BlobMetadata, Literal, Scalar @@ -252,7 +252,7 @@ def to_literal( should_upload = True if python_val is None: - raise AssertionError("None value cannot be converted to a file.") + raise TypeTransformerFailedError("None value cannot be converted to a file.") if not (python_type is os.PathLike or issubclass(python_type, FlyteFile)): raise ValueError(f"Incorrect type {python_type}, must be either a FlyteFile or os.PathLike") @@ -297,13 +297,13 @@ def to_literal( if isinstance(python_val, str): p = pathlib.Path(python_val) if not p.is_file(): - raise ValueError(f"Error converting {python_val} because it's not a file.") + raise TypeTransformerFailedError(f"Error converting {python_val} because it's not a file.") # python_type must be os.PathLike - see check at beginning of function else: should_upload = False else: - raise AssertionError(f"Expected FlyteFile or os.PathLike object, received {type(python_val)}") + raise TypeTransformerFailedError(f"Expected FlyteFile or os.PathLike object, received {type(python_val)}") # If we're uploading something, that means that the uri should always point to the upload destination. if should_upload: diff --git a/flytekit/types/schema/types.py b/flytekit/types/schema/types.py index 89b6da3a90..9864f660b4 100644 --- a/flytekit/types/schema/types.py +++ b/flytekit/types/schema/types.py @@ -11,7 +11,7 @@ import numpy as _np from flytekit.core.context_manager import FlyteContext, FlyteContextManager -from flytekit.core.type_engine import T, TypeEngine, TypeTransformer +from flytekit.core.type_engine import T, TypeEngine, TypeTransformer, TypeTransformerFailedError from flytekit.models.literals import Literal, Scalar, Schema from flytekit.models.types import LiteralType, SchemaType from flytekit.plugins import pandas @@ -361,16 +361,19 @@ def to_literal( local_path=ctx.file_access.get_random_local_directory(), remote_path=ctx.file_access.get_random_remote_directory(), ) + try: + h = SchemaEngine.get_handler(type(python_val)) + except ValueError as e: + raise TypeTransformerFailedError(f"Could not convert {type(python_val)} to flyte schema") from e writer = schema.open(type(python_val)) writer.write(python_val) - h = SchemaEngine.get_handler(type(python_val)) if not h.handles_remote_io: ctx.file_access.put_data(schema.local_path, schema.remote_path, is_multipart=True) return Literal(scalar=Scalar(schema=Schema(schema.remote_path, self._get_schema_type(python_type)))) def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[FlyteSchema]) -> FlyteSchema: if not (lv and lv.scalar and lv.scalar.schema): - raise AssertionError("Can only convert a literal schema to a FlyteSchema") + raise TypeTransformerFailedError("Can only convert a literal schema to a FlyteSchema") def downloader(x, y): ctx.file_access.get_data(x, y, is_multipart=True) diff --git a/tests/flytekit/unit/core/test_type_hints.py b/tests/flytekit/unit/core/test_type_hints.py index 05edf43d6e..e86b6ced05 100644 --- a/tests/flytekit/unit/core/test_type_hints.py +++ b/tests/flytekit/unit/core/test_type_hints.py @@ -7,6 +7,7 @@ import typing from collections import OrderedDict from dataclasses import dataclass +from textwrap import dedent import pandas import pytest @@ -1461,7 +1462,21 @@ def wf2(a: typing.Union[int, str]) -> typing.Union[int, str]: with pytest.raises( TypeError, - match='Cannot convert from scalar {\n primitive {\n string_value: "2"\n }\n}\n to typing.Union\\[float, dict\\]', + match=dedent(r''' + Cannot convert from scalar { + union { + value { + scalar { + primitive { + string_value: "2" + } + } + } + tag: "str" + } + } + to typing.Union\[float, dict\] \(using tag str\) + ''')[1:-1], ): assert wf2(a="2") == "2" From 39b80072a0272192d3cd163bfaef39fa37131f20 Mon Sep 17 00:00:00 2001 From: Maksim Smolin Date: Fri, 10 Dec 2021 11:56:03 -0800 Subject: [PATCH 014/128] Add implicit wrapping union type tests Signed-off-by: maximsmol --- tests/flytekit/unit/core/test_type_hints.py | 26 +++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/tests/flytekit/unit/core/test_type_hints.py b/tests/flytekit/unit/core/test_type_hints.py index e86b6ced05..3d593d9541 100644 --- a/tests/flytekit/unit/core/test_type_hints.py +++ b/tests/flytekit/unit/core/test_type_hints.py @@ -1492,3 +1492,29 @@ def wf(a: typing.Optional[int]) -> typing.Optional[int]: assert wf(a=2) == 2 assert wf(a=None) is None + + +def test_optional_type_implicit_wrapping(): + @task + def t1(a: int) -> typing.Optional[int]: + return a if a > 0 else None + + @workflow + def wf(a: int) -> typing.Optional[int]: + return t1(a=a) + + assert wf(a=2) == 2 + assert wf(a=-10) is None + + +def test_union_type_implicit_wrapping(): + @task + def t1(a: int) -> typing.Union[int, str]: + return a if a > 0 else str(a) + + @workflow + def wf(a: int) -> typing.Union[int, str]: + return t1(a=a) + + assert wf(a=2) == 2 + assert wf(a=-10) == "-10" From 4ff30382bb01269a3ecd82e530eebe6c88e65e5e Mon Sep 17 00:00:00 2001 From: Maksim Smolin Date: Fri, 10 Dec 2021 12:02:42 -0800 Subject: [PATCH 015/128] Add union ambiguity tests Signed-off-by: maximsmol --- tests/flytekit/unit/core/test_type_engine.py | 1 - tests/flytekit/unit/core/test_type_hints.py | 84 +++++++++++++++++++- 2 files changed, 83 insertions(+), 2 deletions(-) diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 5dc053ad6c..933e7dc10c 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -12,7 +12,6 @@ from google.protobuf import struct_pb2 as _struct from marshmallow_jsonschema import JSONSchema -from flytekit.common.exceptions import user as user_exceptions from flytekit.common.types import primitives from flytekit.core.context_manager import FlyteContext, FlyteContextManager from flytekit.core.type_engine import ( diff --git a/tests/flytekit/unit/core/test_type_hints.py b/tests/flytekit/unit/core/test_type_hints.py index 3d593d9541..0a195a5390 100644 --- a/tests/flytekit/unit/core/test_type_hints.py +++ b/tests/flytekit/unit/core/test_type_hints.py @@ -17,6 +17,7 @@ import flytekit from flytekit import ContainerTask, Secret, SQLTask, dynamic, kwtypes, map_task from flytekit.common.translator import get_serializable +from flytekit.common.types import primitives from flytekit.core import context_manager, launch_plan, promise from flytekit.core.condition import conditional from flytekit.core.context_manager import ExecutionState, FastSerializationSettings, Image, ImageConfig @@ -26,7 +27,7 @@ from flytekit.core.resources import Resources from flytekit.core.task import TaskMetadata, task from flytekit.core.testing import patch, task_mock -from flytekit.core.type_engine import RestrictedTypeError, TypeEngine +from flytekit.core.type_engine import RestrictedTypeError, SimpleTransformer, TypeEngine from flytekit.core.workflow import workflow from flytekit.models import literals as _literal_models from flytekit.models.core import types as _core_types @@ -1518,3 +1519,84 @@ def wf(a: int) -> typing.Union[int, str]: assert wf(a=2) == 2 assert wf(a=-10) == "-10" + + +def test_union_type_ambiguity_checking(): + class MyInt: + def __init__(self, x: int): + self.val = x + + def __eq__(self, other): + if not isinstance(other, MyInt): + return False + return other.val == self.val + + TypeEngine.register( + SimpleTransformer( + "MyInt", + MyInt, + primitives.Integer.to_flyte_literal_type(), + lambda x: _literal_models.Literal(scalar=_literal_models.Scalar(primitive=_literal_models.Primitive(integer=x.val))), + lambda x: MyInt(x.scalar.primitive.integer), + ) + ) + + @task + def t1(a: typing.Union[int, MyInt]) -> int: + if isinstance(a, MyInt): + return a.val + return a + + @workflow + def wf(a: int) -> int: + return t1(a=a) + + with pytest.raises( + TypeError, + match="Ambiguous choice of variant for union type. Both int and MyInt transformers match" + ): + assert wf(a=10) == 10 + + del TypeEngine._REGISTRY[MyInt] + + +def test_union_type_ambiguity_resolution(): + class MyInt: + def __init__(self, x: int): + self.val = x + + def __eq__(self, other): + if not isinstance(other, MyInt): + return False + return other.val == self.val + + TypeEngine.register( + SimpleTransformer( + "MyInt", + MyInt, + primitives.Integer.to_flyte_literal_type(), + lambda x: _literal_models.Literal(scalar=_literal_models.Scalar(primitive=_literal_models.Primitive(integer=x.val))), + lambda x: MyInt(x.scalar.primitive.integer), + ) + ) + + @task + def t1(a: typing.Union[int, MyInt]) -> str: + if isinstance(a, MyInt): + return f"MyInt {str(a.val)}" + return str(a) + + @task + def t2(a: int) -> typing.Union[int, MyInt]: + if a < 0: + return MyInt(a) + return a + + @workflow + def wf(a: int) -> str: + return t1(a=t2(a=a)) + + assert wf(a=10) == "10" + assert wf(a=-10) == "MyInt -10" + + del TypeEngine._REGISTRY[MyInt] From de2b7a5870a3d07f86955692a79fd187d39c9dac Mon Sep 17 00:00:00 2001 From: Maksim Smolin Date: Fri, 10 Dec 2021 12:18:28 -0800 Subject: [PATCH 016/128] Fixup tests, make TypeTransformerFailed compatible with all old exception types Signed-off-by: maximsmol --- flytekit/core/type_engine.py | 2 +- tests/flytekit/unit/core/test_schema_types.py | 2 +- tests/flytekit/unit/core/test_type_engine.py | 7 +++---- 3 files changed, 5 insertions(+), 6 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 0720ab1b34..cadc2860d7 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -34,7 +34,7 @@ T = typing.TypeVar("T") DEFINITIONS = "definitions" -class TypeTransformerFailedError(TypeError): +class TypeTransformerFailedError(TypeError, AssertionError, ValueError): ... class TypeTransformer(typing.Generic[T]): diff --git a/tests/flytekit/unit/core/test_schema_types.py b/tests/flytekit/unit/core/test_schema_types.py index 6711ec34de..c151aee345 100644 --- a/tests/flytekit/unit/core/test_schema_types.py +++ b/tests/flytekit/unit/core/test_schema_types.py @@ -24,7 +24,7 @@ def test_assert_type(): schema = FlyteSchema[kwtypes(x=int, y=float)] fst = FlyteSchemaTransformer() lt = fst.get_literal_type(schema) - with pytest.raises(ValueError, match="DataFrames of type are not supported currently"): + with pytest.raises(ValueError, match="Could not convert to flyte schema"): TypeEngine.to_literal(ctx, 3, schema, lt) diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 933e7dc10c..324fde24d1 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -12,6 +12,7 @@ from google.protobuf import struct_pb2 as _struct from marshmallow_jsonschema import JSONSchema +import flytekit.common.exceptions.user as user_exceptions from flytekit.common.types import primitives from flytekit.core.context_manager import FlyteContext, FlyteContextManager from flytekit.core.type_engine import ( @@ -699,10 +700,8 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: assert union_type_tags_unique(lt) ctx = FlyteContextManager.current_context() - try: + with pytest.raises(TypeError, match="Ambiguous choice of variant for union type"): TypeEngine.to_literal(ctx, 3, pt, lt) - except AssertionError as e: - assert str(e) == "Ambiguous choice of variant for union type" del TypeEngine._REGISTRY[UnsignedInt] @@ -831,5 +830,5 @@ def test_dict_to_literal_map_with_wrong_input_type(): ctx = FlyteContext.current_context() input = {"a": 1} guessed_python_types = {"a": str} - with pytest.raises(TypeTransformerFailedError): + with pytest.raises(user_exceptions.FlyteTypeException): TypeEngine.dict_to_literal_map(ctx, input, guessed_python_types) From 596fefcd54ceb1ce182efda5afdef0715f476819 Mon Sep 17 00:00:00 2001 From: Maksim Smolin Date: Fri, 10 Dec 2021 12:28:01 -0800 Subject: [PATCH 017/128] Fixup models + add tests Signed-off-by: maximsmol --- flytekit/models/types.py | 6 ++-- tests/flytekit/common/parameterizers.py | 44 +++++++++++++++++++++++++ 2 files changed, 47 insertions(+), 3 deletions(-) diff --git a/flytekit/models/types.py b/flytekit/models/types.py index 7ba19b10ae..6443d86173 100644 --- a/flytekit/models/types.py +++ b/flytekit/models/types.py @@ -118,7 +118,7 @@ def tag(self) -> str: def to_flyte_idl(self) -> _types_pb2.UnionVariant: return _types_pb2.UnionVariant( - type=self.type, + type=self.type.to_flyte_idl(), tag=self.tag ) @@ -139,8 +139,8 @@ def __init__(self, variants: typing.List["UnionVariant"]): def variants(self) -> typing.List["UnionVariant"]: return self._variants - def to_flyte_idl(self) -> _types_pb2.UnionVariant: - return _types_pb2.UnionVariant( + def to_flyte_idl(self) -> _types_pb2.UnionType: + return _types_pb2.UnionType( variants=[val.to_flyte_idl() if val else None for val in self._variants], ) diff --git a/tests/flytekit/common/parameterizers.py b/tests/flytekit/common/parameterizers.py index fbc8f07ba5..edb1abebd9 100644 --- a/tests/flytekit/common/parameterizers.py +++ b/tests/flytekit/common/parameterizers.py @@ -55,6 +55,20 @@ dimensionality=_core_types.BlobType.BlobDimensionality.MULTIPART, ) ), + types.LiteralType( + union_type=types.UnionType( + variants=[ + types.UnionVariant( + type=types.LiteralType(simple=types.SimpleType.STRING), + tag="str" + ), + types.UnionVariant( + type=types.LiteralType(simple=types.SimpleType.INTEGER), + tag="int" + ) + ] + ) + ), ] @@ -240,6 +254,36 @@ ), ), ), + ( + literals.Scalar( + union=literals.Union( + value=literals.Literal( + scalar=literals.Scalar( + primitive=literals.Primitive( + integer=10 + ) + ) + ), + tag="int" + ) + ), + 10 + ), + ( + literals.Scalar( + union=literals.Union( + value=literals.Literal( + scalar=literals.Scalar( + primitive=literals.Primitive( + string_value="test" + ) + ) + ), + tag="str" + ) + ), + "test" + ) ] LIST_OF_SCALAR_LITERALS_AND_PYTHON_VALUE = [ From 3044e853c4a0655b7dd0d9be66c345a7632a0e03 Mon Sep 17 00:00:00 2001 From: Maksim Smolin Date: Sun, 26 Dec 2021 23:51:17 -0800 Subject: [PATCH 018/128] Implement changed design Signed-off-by: maximsmol --- flytekit/core/type_engine.py | 20 ++++-- flytekit/models/literals.py | 18 +++--- flytekit/models/types.py | 59 +++++++++-------- tests/flytekit/common/parameterizers.py | 20 +++--- tests/flytekit/unit/core/test_type_engine.py | 67 +++++++++++++------- tests/flytekit/unit/core/test_type_hints.py | 7 +- 6 files changed, 113 insertions(+), 78 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index cadc2860d7..b1a4fa2a50 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -29,7 +29,7 @@ from flytekit.models import types as _type_models from flytekit.models.core import types as _core_types from flytekit.models.literals import Literal, LiteralCollection, LiteralMap, Primitive, Scalar, Union, Void -from flytekit.models.types import LiteralType, SimpleType, UnionType, UnionVariant +from flytekit.models.types import LiteralType, SimpleType, TypeStructure, UnionType T = typing.TypeVar("T") DEFINITIONS = "definitions" @@ -617,6 +617,9 @@ def guess_python_type(self, literal_type: LiteralType) -> Type[list]: return typing.List[ct] raise ValueError(f"List transformer cannot reverse {literal_type}") +def _add_tag_to_type(x: LiteralType, tag: str) -> LiteralType: + x._structure = TypeStructure(tag=tag) + return x class UnionTransformer(TypeTransformer[T]): """ @@ -627,9 +630,10 @@ def __init__(self): super().__init__("Typed Union", typing.Union) def get_literal_type(self, t: Type[T]) -> Optional[LiteralType]: + try: trans = [(TypeEngine.get_transformer(x), x) for x in typing.get_args(t)] - variants = [UnionVariant(t.get_literal_type(x), t.name) for (t, x) in trans] + variants = [_add_tag_to_type(t.get_literal_type(x), t.name) for (t, x) in trans] return _type_models.LiteralType(union_type=UnionType(variants)) except Exception as e: raise ValueError(f"Type of Generic Union type is not supported, {e}") @@ -637,13 +641,13 @@ def get_literal_type(self, t: Type[T]) -> Optional[LiteralType]: def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], expected: LiteralType) -> Literal: found_res = False res = None - res_tag = None + res_type = None for t in typing.get_args(python_type): try: trans = TypeEngine.get_transformer(t) res = trans.to_literal(ctx, python_val, t, expected) - res_tag = trans.name + res_type = _add_tag_to_type(trans.get_literal_type(t), trans.name) if found_res: # Should really never happen, sanity check raise TypeError(f"Ambiguous choice of variant for union type") @@ -653,14 +657,16 @@ def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], exp continue if found_res: - return Literal(scalar=Scalar(union=Union(value=res, tag=res_tag))) + return Literal(scalar=Scalar(union=Union(value=res, stored_type=res_type))) raise TypeTransformerFailedError(f"Cannot convert from {python_val} to {python_type}") def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[T]) -> Optional[typing.Any]: union_tag = None if lv.scalar is not None and lv.scalar.union is not None: - union_tag = lv.scalar.union.tag + union_type = lv.scalar.union.stored_type + if union_type.structure is not None: + union_tag = union_type.structure.tag found_res = False res = None @@ -668,7 +674,7 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: for v in typing.get_args(expected_python_type): try: trans = TypeEngine.get_transformer(v) - if union_tag: + if union_tag is not None: if trans.name != union_tag: continue diff --git a/flytekit/models/literals.py b/flytekit/models/literals.py index 166b03e6e6..b3acc7abde 100644 --- a/flytekit/models/literals.py +++ b/flytekit/models/literals.py @@ -7,9 +7,9 @@ from flytekit.common.exceptions import user as _user_exceptions from flytekit.models import common as _common from flytekit.models.core import types as _core_types +from flytekit.models.types import LiteralType as _LiteralType from flytekit.models.types import OutputReference as _OutputReference from flytekit.models.types import SchemaType as _SchemaType -from flytekit.models.types import UnionType as _UnionType class RetryStrategy(_common.FlyteIdlEntity): @@ -548,15 +548,15 @@ def from_flyte_idl(cls, pb2_object): class Union(_common.FlyteIdlEntity): - def __init__(self, value, tag): + def __init__(self, value, stored_type): """ The runtime representation of a tagged union value. See `UnionType` for more details. :param flytekit.models.literals.Literal value: - :param str tag: + :param flytekit.models.types.LiteralType stored_type: """ self._value = value - self._tag = tag + self._type = stored_type @property def value(self): @@ -566,17 +566,17 @@ def value(self): return self._value @property - def tag(self): + def stored_type(self): """ - :rtype: str + :rtype: flytekit.models.types.LiteralType """ - return self._tag + return self._type def to_flyte_idl(self): """ :rtype: flyteidl.core.literals_pb2.Union """ - return _literals_pb2.Union(value=self.value.to_flyte_idl(), tag=self.tag) + return _literals_pb2.Union(value=self.value.to_flyte_idl(), type=self._type.to_flyte_idl()) @classmethod def from_flyte_idl(cls, pb2_object): @@ -584,7 +584,7 @@ def from_flyte_idl(cls, pb2_object): :param flyteidl.core.literals_pb2.Schema pb2_object: :rtype: Schema """ - return cls(value=Literal.from_flyte_idl(pb2_object.value), tag=pb2_object.tag) + return cls(value=Literal.from_flyte_idl(pb2_object.value), stored_type=_LiteralType.from_flyte_idl(pb2_object.type)) class LiteralCollection(_common.FlyteIdlEntity): def __init__(self, literals): diff --git a/flytekit/models/types.py b/flytekit/models/types.py index 6443d86173..d967e1bfb4 100644 --- a/flytekit/models/types.py +++ b/flytekit/models/types.py @@ -99,54 +99,48 @@ def from_flyte_idl(cls, proto): return cls(columns=[SchemaType.SchemaColumn.from_flyte_idl(c) for c in proto.columns]) -class UnionVariant(_common.FlyteIdlEntity): +class UnionType(_common.FlyteIdlEntity): """ - Models _types_pb2.UnionVariant + Models _types_pb2.UnionType """ - def __init__(self, type: "LiteralType", tag: str): - self._type = type - self._tag = tag - - @property - def type(self) -> "LiteralType": - return self._type + def __init__(self, variants: typing.List["LiteralType"]): + self._variants = variants @property - def tag(self) -> str: - return self._tag + def variants(self) -> typing.List["LiteralType"]: + return self._variants - def to_flyte_idl(self) -> _types_pb2.UnionVariant: - return _types_pb2.UnionVariant( - type=self.type.to_flyte_idl(), - tag=self.tag + def to_flyte_idl(self) -> _types_pb2.UnionType: + return _types_pb2.UnionType( + variants=[val.to_flyte_idl() if val else None for val in self._variants], ) @classmethod - def from_flyte_idl(cls, proto: _types_pb2.UnionVariant): - return cls(type=LiteralType.from_flyte_idl(proto.type), tag=proto.tag) + def from_flyte_idl(cls, proto: _types_pb2.UnionType): + return cls(variants=[LiteralType.from_flyte_idl(v) for v in proto.variants]) -class UnionType(_common.FlyteIdlEntity): +class TypeStructure(_common.FlyteIdlEntity): """ - Models _types_pb2.UnionType + Models _types_pb2.TypeStructure """ - def __init__(self, variants: typing.List["UnionVariant"]): - self._variants = variants + def __init__(self, tag: str): + self._tag = tag @property - def variants(self) -> typing.List["UnionVariant"]: - return self._variants + def tag(self) -> str: + return self._tag - def to_flyte_idl(self) -> _types_pb2.UnionType: - return _types_pb2.UnionType( - variants=[val.to_flyte_idl() if val else None for val in self._variants], + def to_flyte_idl(self) -> _types_pb2.TypeStructure: + return _types_pb2.TypeStructure( + tag=self._tag, ) @classmethod - def from_flyte_idl(cls, proto: _types_pb2.UnionType): - return cls(variants=[UnionVariant.from_flyte_idl(v) for v in proto.variants]) + def from_flyte_idl(cls, proto: _types_pb2.TypeStructure): + return cls(tag=proto.tag) class LiteralType(_common.FlyteIdlEntity): @@ -159,6 +153,7 @@ def __init__( blob=None, enum_type=None, union_type=None, + structure=None, metadata=None, ): """ @@ -172,6 +167,7 @@ def __init__( :param flytekit.models.core.types.BlobType blob: For blob objects, this describes the type. :param flytekit.models.core.types.EnumType enum_type: For enum objects, describes an enum :param flytekit.models.core.types.UnionType union_type: For union objects, describes an python union type. + :param flytekit.models.core.types.TypeStructure structure: Type matching hints :param dict[Text, T] metadata: Additional data describing the type """ self._simple = simple @@ -181,6 +177,7 @@ def __init__( self._blob = blob self._enum_type = enum_type self._union_type = union_type + self._structure = structure self._metadata = metadata @property @@ -217,6 +214,10 @@ def enum_type(self) -> _core_types.EnumType: def union_type(self) -> UnionType: return self._union_type + @property + def structure(self) -> TypeStructure: + return self._structure + @property def metadata(self): """ @@ -244,6 +245,7 @@ def to_flyte_idl(self): blob=self.blob.to_flyte_idl() if self.blob is not None else None, enum_type=self.enum_type.to_flyte_idl() if self.enum_type else None, union_type=self.union_type.to_flyte_idl() if self.union_type else None, + structure=self.structure.to_flyte_idl() if self.structure else None, metadata=metadata, ) return t @@ -268,6 +270,7 @@ def from_flyte_idl(cls, proto): blob=_core_types.BlobType.from_flyte_idl(proto.blob) if proto.HasField("blob") else None, enum_type=_core_types.EnumType.from_flyte_idl(proto.enum_type) if proto.HasField("enum_type") else None, union_type=UnionType.from_flyte_idl(proto.union_type) if proto.HasField("union_type") else None, + structure=TypeStructure.from_flyte_idl(proto.structure) if proto.HasField("structure") else None, metadata=_json_format.MessageToDict(proto.metadata) or None, ) diff --git a/tests/flytekit/common/parameterizers.py b/tests/flytekit/common/parameterizers.py index edb1abebd9..c7d022d860 100644 --- a/tests/flytekit/common/parameterizers.py +++ b/tests/flytekit/common/parameterizers.py @@ -58,14 +58,8 @@ types.LiteralType( union_type=types.UnionType( variants=[ - types.UnionVariant( - type=types.LiteralType(simple=types.SimpleType.STRING), - tag="str" - ), - types.UnionVariant( - type=types.LiteralType(simple=types.SimpleType.INTEGER), - tag="int" - ) + types.LiteralType(simple=types.SimpleType.STRING, structure=types.TypeStructure(tag="str")), + types.LiteralType(simple=types.SimpleType.INTEGER, structure=types.TypeStructure(tag="int")), ] ) ), @@ -264,7 +258,10 @@ ) ) ), - tag="int" + stored_type=types.LiteralType( + simple=types.SimpleType.INTEGER, + structure=types.TypeStructure(tag="int") + ) ) ), 10 @@ -279,7 +276,10 @@ ) ) ), - tag="str" + stored_type=types.LiteralType( + simple=types.SimpleType.STRING, + structure=types.TypeStructure(tag="str") + ) ) ), "test" diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 324fde24d1..da2d5f6dcf 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -29,7 +29,7 @@ from flytekit.models import types as model_types from flytekit.models.core.types import BlobType from flytekit.models.literals import Blob, BlobMetadata, Literal, LiteralCollection, LiteralMap, Primitive, Scalar, Void -from flytekit.models.types import LiteralType, SimpleType +from flytekit.models.types import LiteralType, SimpleType, TypeStructure from flytekit.types.directory.types import FlyteDirectory from flytekit.types.file import JPEGImageFile from flytekit.types.file.file import FlyteFile, FlyteFilePathTransformer @@ -554,9 +554,9 @@ def test_enum_type(): def union_type_tags_unique(t: LiteralType): seen = set() for x in t.union_type.variants: - if x.tag in seen: + if x.structure.tag in seen: return False - seen.add(x.tag) + seen.add(x.structure.tag) return True @@ -564,19 +564,22 @@ def union_type_tags_unique(t: LiteralType): def test_union_type(): pt = typing.Union[str, int] lt = TypeEngine.to_literal_type(pt) - assert [x.type for x in lt.union_type.variants] == [LiteralType(simple=SimpleType.STRING), LiteralType(simple=SimpleType.INTEGER)] + assert lt.union_type.variants == [ + LiteralType(simple=SimpleType.STRING, structure=TypeStructure(tag="str")), + LiteralType(simple=SimpleType.INTEGER, structure=TypeStructure(tag="int")) + ] assert union_type_tags_unique(lt) ctx = FlyteContextManager.current_context() lv = TypeEngine.to_literal(ctx, 3, pt, lt) v = TypeEngine.to_python_value(ctx, lv, pt) - assert lv.scalar.union.tag == "int" + assert lv.scalar.union.stored_type.structure.tag == "int" assert lv.scalar.union.value.scalar.primitive.integer == 3 assert v == 3 lv = TypeEngine.to_literal(ctx, "hello", pt, lt) v = TypeEngine.to_python_value(ctx, lv, pt) - assert lv.scalar.union.tag == "str" + assert lv.scalar.union.stored_type.structure.tag == "str" assert lv.scalar.union.value.scalar.primitive.string_value == "hello" assert v == "hello" @@ -584,26 +587,32 @@ def test_union_type(): def test_optional_type(): pt = typing.Optional[int] lt = TypeEngine.to_literal_type(pt) - assert [x.type for x in lt.union_type.variants] == [LiteralType(simple=SimpleType.INTEGER), LiteralType(simple=SimpleType.NONE)] + assert lt.union_type.variants == [ + LiteralType(simple=SimpleType.INTEGER, structure=TypeStructure(tag="int")), + LiteralType(simple=SimpleType.NONE, structure=TypeStructure(tag="none")) + ] assert union_type_tags_unique(lt) ctx = FlyteContextManager.current_context() lv = TypeEngine.to_literal(ctx, 3, pt, lt) v = TypeEngine.to_python_value(ctx, lv, pt) - assert lv.scalar.union.tag == "int" + assert lv.scalar.union.stored_type.structure.tag == "int" assert lv.scalar.union.value.scalar.primitive.integer == 3 assert v == 3 lv = TypeEngine.to_literal(ctx, None, pt, lt) v = TypeEngine.to_python_value(ctx, lv, pt) - assert lv.scalar.union.tag == "none" + assert lv.scalar.union.stored_type.structure.tag == "none" assert lv.scalar.union.value.scalar.none_type == Void() assert v is None def test_union_from_unambiguous_literal(): pt = typing.Union[str, int] lt = TypeEngine.to_literal_type(pt) - assert [x.type for x in lt.union_type.variants] == [LiteralType(simple=SimpleType.STRING), LiteralType(simple=SimpleType.INTEGER)] + assert lt.union_type.variants == [ + LiteralType(simple=SimpleType.STRING, structure=TypeStructure(tag="str")), + LiteralType(simple=SimpleType.INTEGER, structure=TypeStructure(tag="int")) + ] assert union_type_tags_unique(lt) ctx = FlyteContextManager.current_context() @@ -635,19 +644,22 @@ def __eq__(self, other): pt = typing.Union[int, MyInt] lt = TypeEngine.to_literal_type(pt) - assert [x.type for x in lt.union_type.variants] == [LiteralType(simple=SimpleType.INTEGER), LiteralType(simple=SimpleType.INTEGER)] + assert lt.union_type.variants == [ + LiteralType(simple=SimpleType.INTEGER, structure=TypeStructure(tag="int")), + LiteralType(simple=SimpleType.INTEGER, structure=TypeStructure(tag="MyInt")) + ] assert union_type_tags_unique(lt) ctx = FlyteContextManager.current_context() lv = TypeEngine.to_literal(ctx, 3, pt, lt) v = TypeEngine.to_python_value(ctx, lv, pt) - assert lv.scalar.union.tag == "int" + assert lv.scalar.union.stored_type.structure.tag == "int" assert lv.scalar.union.value.scalar.primitive.integer == 3 assert v == 3 lv = TypeEngine.to_literal(ctx, MyInt(10), pt, lt) v = TypeEngine.to_python_value(ctx, lv, pt) - assert lv.scalar.union.tag == "MyInt" + assert lv.scalar.union.stored_type.structure.tag == "MyInt" assert lv.scalar.union.value.scalar.primitive.integer == 10 assert v == MyInt(10) @@ -696,7 +708,10 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: pt = typing.Union[int, UnsignedInt] lt = TypeEngine.to_literal_type(pt) - assert [x.type for x in lt.union_type.variants] == [LiteralType(simple=SimpleType.INTEGER), LiteralType(simple=SimpleType.INTEGER)] + assert lt.union_type.variants == [ + LiteralType(simple=SimpleType.INTEGER, structure=TypeStructure(tag="int")), + LiteralType(simple=SimpleType.INTEGER, structure=TypeStructure(tag="UnsignedInt")) + ] assert union_type_tags_unique(lt) ctx = FlyteContextManager.current_context() @@ -709,22 +724,28 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: def test_union_of_lists(): pt = typing.Union[typing.List[int], typing.List[str]] lt = TypeEngine.to_literal_type(pt) - assert [x.type for x in lt.union_type.variants] == [ - LiteralType(collection_type=LiteralType(simple=SimpleType.INTEGER)), - LiteralType(collection_type=LiteralType(simple=SimpleType.STRING)), + assert lt.union_type.variants == [ + LiteralType( + collection_type=LiteralType(simple=SimpleType.INTEGER, structure=TypeStructure(tag="int")), + structure=TypeStructure(tag="Typed List") + ), + LiteralType( + collection_type=LiteralType(simple=SimpleType.STRING, structure=TypeStructure(tag="str")), + structure=TypeStructure(tag="Typed List") + ), ] assert not union_type_tags_unique(lt) # tags are deliberately NOT unique ctx = FlyteContextManager.current_context() lv = TypeEngine.to_literal(ctx, ["hello", "world"], pt, lt) v = TypeEngine.to_python_value(ctx, lv, pt) - assert lv.scalar.union.tag == "Typed List" + assert lv.scalar.union.stored_type.structure.tag == "Typed List" assert [x.scalar.primitive.string_value for x in lv.scalar.union.value.collection.literals] == ["hello", "world"] assert v == ["hello", "world"] lv = TypeEngine.to_literal(ctx, [1, 3], pt, lt) v = TypeEngine.to_python_value(ctx, lv, pt) - assert lv.scalar.union.tag == "Typed List" + assert lv.scalar.union.stored_type.structure.tag == "Typed List" assert [x.scalar.primitive.integer for x in lv.scalar.union.value.collection.literals] == [1, 3] assert v == [1, 3] @@ -733,16 +754,16 @@ def test_list_of_unions(): pt = typing.List[typing.Union[str, int]] lt = TypeEngine.to_literal_type(pt) # todo(maximsmol): seems like the order here is non-deterministic - assert [x.type for x in lt.collection_type.union_type.variants] == [ - LiteralType(simple=SimpleType.STRING), - LiteralType(simple=SimpleType.INTEGER), + assert lt.collection_type.union_type.variants == [ + LiteralType(simple=SimpleType.STRING, structure=TypeStructure(tag="str")), + LiteralType(simple=SimpleType.INTEGER, structure=TypeStructure(tag="int")), ] assert union_type_tags_unique(lt.collection_type) # tags are deliberately NOT unique ctx = FlyteContextManager.current_context() lv = TypeEngine.to_literal(ctx, ["hello", 123, "world"], pt, lt) v = TypeEngine.to_python_value(ctx, lv, pt) - assert [x.scalar.union.tag for x in lv.collection.literals] == ["str", "int", "str"] + assert [x.scalar.union.stored_type.structure.tag for x in lv.collection.literals] == ["str", "int", "str"] assert v == ["hello", 123, "world"] @pytest.mark.parametrize( diff --git a/tests/flytekit/unit/core/test_type_hints.py b/tests/flytekit/unit/core/test_type_hints.py index 0a195a5390..1e1f88a756 100644 --- a/tests/flytekit/unit/core/test_type_hints.py +++ b/tests/flytekit/unit/core/test_type_hints.py @@ -1473,7 +1473,12 @@ def wf2(a: typing.Union[int, str]) -> typing.Union[int, str]: } } } - tag: "str" + type { + simple: STRING + structure { + tag: "str" + } + } } } to typing.Union\[float, dict\] \(using tag str\) From a9e3c404c8b7f2b986364c9e5bfe54427750e853 Mon Sep 17 00:00:00 2001 From: Maksim Smolin Date: Thu, 17 Feb 2022 08:29:06 -0800 Subject: [PATCH 019/128] Fix PR comments Signed-off-by: maximsmol --- flytekit/core/type_engine.py | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index b1a4fa2a50..6dcd84e93f 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -31,6 +31,15 @@ from flytekit.models.literals import Literal, LiteralCollection, LiteralMap, Primitive, Scalar, Union, Void from flytekit.models.types import LiteralType, SimpleType, TypeStructure, UnionType +try: + from typing import get_args as _get_args +except ImportError: + try: + from typing_extensions import get_args as _get_args + except ImportError: + def _get_args(t): + return t.__args__ + T = typing.TypeVar("T") DEFINITIONS = "definitions" @@ -623,16 +632,15 @@ def _add_tag_to_type(x: LiteralType, tag: str) -> LiteralType: class UnionTransformer(TypeTransformer[T]): """ - Transformer that handles a univariate typing.Union[T] + Transformer that handles a typing.Union[T1, T2, ...] """ def __init__(self): super().__init__("Typed Union", typing.Union) def get_literal_type(self, t: Type[T]) -> Optional[LiteralType]: - try: - trans = [(TypeEngine.get_transformer(x), x) for x in typing.get_args(t)] + trans = [(TypeEngine.get_transformer(x), x) for x in _get_args(t)] variants = [_add_tag_to_type(t.get_literal_type(x), t.name) for (t, x) in trans] return _type_models.LiteralType(union_type=UnionType(variants)) except Exception as e: @@ -642,7 +650,7 @@ def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], exp found_res = False res = None res_type = None - for t in typing.get_args(python_type): + for t in _get_args(python_type): try: trans = TypeEngine.get_transformer(t) @@ -671,7 +679,7 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: found_res = False res = None res_tag = None - for v in typing.get_args(expected_python_type): + for v in _get_args(expected_python_type): try: trans = TypeEngine.get_transformer(v) if union_tag is not None: @@ -696,7 +704,7 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: f"Both {res_tag} and {trans.name} transformers match") res_tag = trans.name found_res = True - except TypeTransformerFailedError as e: # todo(maximsmol): separate programmer errors from type conversion failures + except TypeTransformerFailedError as e: logger.debug(f"Failed to convert from {lv} to {v}", e) if found_res: @@ -887,7 +895,7 @@ def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], exp if type(python_val).__class__ != enum.EnumMeta: raise TypeTransformerFailedError("Expected an enum") if type(python_val.value) != str: - raise TypeTransformerFailedError("Only string-valued enums are sujpportedd") + raise TypeTransformerFailedError("Only string-valued enums are supportedd") return Literal(scalar=Scalar(primitive=Primitive(string_value=python_val.value))) # type: ignore From f339931437161af74e20aa1172949cdcf449bca1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Emirhan=20Karag=C3=BCl?= Date: Wed, 17 Nov 2021 22:28:06 +0100 Subject: [PATCH 020/128] Remote entrypoint serialize (#733) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Emirhan Karagül Signed-off-by: Yee Hing Tong Signed-off-by: maximsmol --- flytekit/remote/remote.py | 25 +++++++++++- .../tests/test_remote_register.py | 40 +++++++++++++++++++ 2 files changed, 63 insertions(+), 2 deletions(-) create mode 100644 plugins/flytekit-spark/tests/test_remote_register.py diff --git a/flytekit/remote/remote.py b/flytekit/remote/remote.py index 4ba7032437..dccf62505f 100644 --- a/flytekit/remote/remote.py +++ b/flytekit/remote/remote.py @@ -34,13 +34,19 @@ from flytekit.clients.helpers import iterate_node_executions, iterate_task_executions from flytekit.clis.flyte_cli.main import _detect_default_config_file +from flytekit.clis.sdk_in_container import serialize from flytekit.common import constants from flytekit.common.exceptions import user as user_exceptions from flytekit.common.translator import FlyteControlPlaneEntity, FlyteLocalEntity, get_serializable from flytekit.configuration import auth as auth_config from flytekit.configuration.internal import DOMAIN, PROJECT from flytekit.core.base_task import PythonTask -from flytekit.core.context_manager import FlyteContextManager, ImageConfig, SerializationSettings, get_image_config +from flytekit.core.context_manager import ( + FlyteContextManager, + ImageConfig, + SerializationSettings, + get_image_config, +) from flytekit.core.data_persistence import FileAccessProvider from flytekit.core.launch_plan import LaunchPlan from flytekit.core.type_engine import TypeEngine @@ -128,6 +134,7 @@ def from_config( default_domain: typing.Optional[str] = None, config_file_path: typing.Optional[str] = None, grpc_credentials: typing.Optional[grpc.ChannelCredentials] = None, + venv_root: typing.Optional[str] = None, ) -> FlyteRemote: """Create a FlyteRemote object using flyte configuration variables and/or environment variable overrides. @@ -151,6 +158,11 @@ def from_config( raw_output_prefix=raw_output_data_prefix, ) + venv_root = venv_root or serialize._DEFAULT_FLYTEKIT_VIRTUALENV_ROOT + entrypoint = context_manager.EntrypointSettings( + path=os.path.join(venv_root, serialize._DEFAULT_FLYTEKIT_RELATIVE_ENTRYPOINT_LOC) + ) + return cls( flyte_admin_url=platform_config.URL.get(), insecure=platform_config.INSECURE.get(), @@ -169,6 +181,7 @@ def from_config( common_models.RawOutputDataConfig(raw_output_data_prefix) if raw_output_data_prefix else None ), grpc_credentials=grpc_credentials, + entrypoint_settings=entrypoint, ) def __init__( @@ -185,6 +198,7 @@ def __init__( image_config: typing.Optional[ImageConfig] = None, raw_output_data_config: typing.Optional[common_models.RawOutputDataConfig] = None, grpc_credentials: typing.Optional[grpc.ChannelCredentials] = None, + entrypoint_settings: typing.Optional[context_manager.EntrypointSettings] = None, ): """Initialize a FlyteRemote object. @@ -199,7 +213,11 @@ def __init__( :param annotations: annotation config :param image_config: image config :param raw_output_data_config: location for offloaded data, e.g. in S3 - :param grpc_credentials: gRPC channel credentials for connecting to flyte admin as returned by :func:`grpc.ssl_channel_credentials` + :param grpc_credentials: gRPC channel credentials for connecting to flyte admin as returned + by :func:`grpc.ssl_channel_credentials` + :param entrypoint_settings: EntrypointSettings object for use with Spark tasks. If supplied, this will be + used when serializing Spark tasks, which need to know the path to the flytekit entrypoint.py file, + inside the container. """ remote_logger.warning("This feature is still in beta. Its interface and UX is subject to change.") if flyte_admin_url is None: @@ -217,6 +235,8 @@ def __init__( self._labels = labels self._annotations = annotations self._raw_output_data_config = raw_output_data_config + # Not exposing this as a property for now. + self._entrypoint_settings = entrypoint_settings # Save the file access object locally, but also make it available for use from the context. FlyteContextManager.with_context(FlyteContextManager.current_context().with_file_access(file_access).build()) @@ -525,6 +545,7 @@ def _serialize( self.image_config, # https://github.com/flyteorg/flyte/issues/1359 env={internal.IMAGE.env_var: self.image_config.default_image.full}, + entrypoint_settings=self._entrypoint_settings, ), entity=entity, ) diff --git a/plugins/flytekit-spark/tests/test_remote_register.py b/plugins/flytekit-spark/tests/test_remote_register.py new file mode 100644 index 0000000000..67d0f63b1f --- /dev/null +++ b/plugins/flytekit-spark/tests/test_remote_register.py @@ -0,0 +1,40 @@ +from flytekitplugins.spark import Spark +from mock import MagicMock, patch + +from flytekit import task +from flytekit.remote.remote import FlyteRemote + + +@patch("flytekit.configuration.platform.URL") +@patch("flytekit.configuration.platform.INSECURE") +def test_spark_template_with_remote(mock_insecure, mock_url): + @task(task_config=Spark(spark_conf={"spark": "1"})) + def my_spark(a: str) -> int: + return 10 + + @task + def my_python_task(a: str) -> int: + return 10 + + mock_url.get.return_value = "localhost" + + mock_insecure.get.return_value = True + mock_client = MagicMock() + + remote = FlyteRemote.from_config("p1", "d1") + + remote._image_config = MagicMock() + remote._client = mock_client + + remote.register(my_spark) + serialized_spec = mock_client.create_task.call_args.kwargs["task_spec"] + + # Check if the serialized spark task has mainApplicaitonFile field set. + assert serialized_spec.template.custom["mainApplicationFile"] + assert serialized_spec.template.custom["sparkConf"] + + remote.register(my_python_task) + serialized_spec = mock_client.create_task.call_args.kwargs["task_spec"] + + # Check if the serialized python task has no mainApplicaitonFile field set by default. + assert serialized_spec.template.custom is None From 53975f04c03ccabf5f9179291a6a1802626e6c4c Mon Sep 17 00:00:00 2001 From: Eduardo Apolinario <653394+eapolinario@users.noreply.github.com> Date: Wed, 17 Nov 2021 13:47:16 -0800 Subject: [PATCH 021/128] Fix lint error in remote.py (#755) Signed-off-by: Eduardo Apolinario Co-authored-by: Eduardo Apolinario Signed-off-by: maximsmol --- flytekit/remote/remote.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/flytekit/remote/remote.py b/flytekit/remote/remote.py index dccf62505f..e859db730d 100644 --- a/flytekit/remote/remote.py +++ b/flytekit/remote/remote.py @@ -41,12 +41,7 @@ from flytekit.configuration import auth as auth_config from flytekit.configuration.internal import DOMAIN, PROJECT from flytekit.core.base_task import PythonTask -from flytekit.core.context_manager import ( - FlyteContextManager, - ImageConfig, - SerializationSettings, - get_image_config, -) +from flytekit.core.context_manager import FlyteContextManager, ImageConfig, SerializationSettings, get_image_config from flytekit.core.data_persistence import FileAccessProvider from flytekit.core.launch_plan import LaunchPlan from flytekit.core.type_engine import TypeEngine From 06254e3d279fd87cfb4f3407e875df5efafa8aa2 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Wed, 24 Nov 2021 03:45:42 +0800 Subject: [PATCH 022/128] Support enum in dataclass (#753) * Add support enum in dataclass Signed-off-by: Kevin Su * Update test Signed-off-by: Kevin Su * Fixed lint Signed-off-by: Kevin Su --- flytekit/core/type_engine.py | 9 +++- tests/flytekit/unit/core/test_type_engine.py | 51 +++++++++++++++----- tests/flytekit/unit/core/test_type_hints.py | 41 +++++++++++++--- 3 files changed, 83 insertions(+), 18 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 6dcd84e93f..da8cfe15f9 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -18,6 +18,7 @@ from google.protobuf.json_format import MessageToDict as _MessageToDict from google.protobuf.json_format import ParseDict as _ParseDict from google.protobuf.struct_pb2 import Struct +from marshmallow_enum import EnumField, LoadDumpOptions from marshmallow_jsonschema import JSONSchema from flytekit.common.exceptions import user as user_exceptions @@ -251,7 +252,13 @@ def get_literal_type(self, t: Type[T]) -> LiteralType: ) schema = None try: - schema = JSONSchema().dump(cast(DataClassJsonMixin, t).schema()) + s = cast(DataClassJsonMixin, t).schema() + for _, v in s.fields.items(): + # marshmallow-jsonschema only supports enums loaded by name. + # https://github.com/fuhrysteve/marshmallow-jsonschema/blob/81eada1a0c42ff67de216923968af0a6b54e5dcb/marshmallow_jsonschema/base.py#L228 + if isinstance(v, EnumField): + v.load_by = LoadDumpOptions.name + schema = JSONSchema().dump(s) except Exception as e: logger.warn("failed to extract schema for object %s, (will run schemaless) error: %s", str(t), e) diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index da2d5f6dcf..1958020065 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -10,6 +10,7 @@ from flyteidl.core import errors_pb2 from google.protobuf import json_format as _json_format from google.protobuf import struct_pb2 as _struct +from marshmallow_enum import LoadDumpOptions from marshmallow_jsonschema import JSONSchema import flytekit.common.exceptions.user as user_exceptions @@ -38,6 +39,7 @@ T = typing.TypeVar("T") + def test_type_engine(): t = int lt = TypeEngine.to_literal_type(t) @@ -566,7 +568,7 @@ def test_union_type(): lt = TypeEngine.to_literal_type(pt) assert lt.union_type.variants == [ LiteralType(simple=SimpleType.STRING, structure=TypeStructure(tag="str")), - LiteralType(simple=SimpleType.INTEGER, structure=TypeStructure(tag="int")) + LiteralType(simple=SimpleType.INTEGER, structure=TypeStructure(tag="int")), ] assert union_type_tags_unique(lt) @@ -589,7 +591,7 @@ def test_optional_type(): lt = TypeEngine.to_literal_type(pt) assert lt.union_type.variants == [ LiteralType(simple=SimpleType.INTEGER, structure=TypeStructure(tag="int")), - LiteralType(simple=SimpleType.NONE, structure=TypeStructure(tag="none")) + LiteralType(simple=SimpleType.NONE, structure=TypeStructure(tag="none")), ] assert union_type_tags_unique(lt) @@ -606,12 +608,13 @@ def test_optional_type(): assert lv.scalar.union.value.scalar.none_type == Void() assert v is None + def test_union_from_unambiguous_literal(): pt = typing.Union[str, int] lt = TypeEngine.to_literal_type(pt) assert lt.union_type.variants == [ LiteralType(simple=SimpleType.STRING, structure=TypeStructure(tag="str")), - LiteralType(simple=SimpleType.INTEGER, structure=TypeStructure(tag="int")) + LiteralType(simple=SimpleType.INTEGER, structure=TypeStructure(tag="int")), ] assert union_type_tags_unique(lt) @@ -622,6 +625,7 @@ def test_union_from_unambiguous_literal(): v = TypeEngine.to_python_value(ctx, lv, pt) assert v == 3 + def test_union_custom_transformer(): class MyInt: def __init__(self, x: int): @@ -646,7 +650,7 @@ def __eq__(self, other): lt = TypeEngine.to_literal_type(pt) assert lt.union_type.variants == [ LiteralType(simple=SimpleType.INTEGER, structure=TypeStructure(tag="int")), - LiteralType(simple=SimpleType.INTEGER, structure=TypeStructure(tag="MyInt")) + LiteralType(simple=SimpleType.INTEGER, structure=TypeStructure(tag="MyInt")), ] assert union_type_tags_unique(lt) @@ -672,6 +676,7 @@ def __eq__(self, other): del TypeEngine._REGISTRY[MyInt] + def test_union_custom_transformer_sanity_check(): class UnsignedInt: def __init__(self, x: int): @@ -690,7 +695,9 @@ def __init__(self): def get_literal_type(self, t: typing.Type[T]) -> LiteralType: return primitives.Integer.to_flyte_literal_type() - def to_literal(self, ctx: FlyteContext, python_val: T, python_type: typing.Type[T], expected: LiteralType) -> Literal: + def to_literal( + self, ctx: FlyteContext, python_val: T, python_type: typing.Type[T], expected: LiteralType + ) -> Literal: if type(python_val) != int: raise TypeTransformerFailedError("Expected an integer") @@ -705,12 +712,11 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: TypeEngine.register(UnsignedIntTransformer()) - pt = typing.Union[int, UnsignedInt] lt = TypeEngine.to_literal_type(pt) assert lt.union_type.variants == [ LiteralType(simple=SimpleType.INTEGER, structure=TypeStructure(tag="int")), - LiteralType(simple=SimpleType.INTEGER, structure=TypeStructure(tag="UnsignedInt")) + LiteralType(simple=SimpleType.INTEGER, structure=TypeStructure(tag="UnsignedInt")), ] assert union_type_tags_unique(lt) @@ -727,14 +733,14 @@ def test_union_of_lists(): assert lt.union_type.variants == [ LiteralType( collection_type=LiteralType(simple=SimpleType.INTEGER, structure=TypeStructure(tag="int")), - structure=TypeStructure(tag="Typed List") + structure=TypeStructure(tag="Typed List"), ), LiteralType( collection_type=LiteralType(simple=SimpleType.STRING, structure=TypeStructure(tag="str")), - structure=TypeStructure(tag="Typed List") + structure=TypeStructure(tag="Typed List"), ), ] - assert not union_type_tags_unique(lt) # tags are deliberately NOT unique + assert not union_type_tags_unique(lt) # tags are deliberately NOT unique ctx = FlyteContextManager.current_context() lv = TypeEngine.to_literal(ctx, ["hello", "world"], pt, lt) @@ -758,7 +764,7 @@ def test_list_of_unions(): LiteralType(simple=SimpleType.STRING, structure=TypeStructure(tag="str")), LiteralType(simple=SimpleType.INTEGER, structure=TypeStructure(tag="int")), ] - assert union_type_tags_unique(lt.collection_type) # tags are deliberately NOT unique + assert union_type_tags_unique(lt.collection_type) # tags are deliberately NOT unique ctx = FlyteContextManager.current_context() lv = TypeEngine.to_literal(ctx, ["hello", 123, "world"], pt, lt) @@ -766,6 +772,29 @@ def test_list_of_unions(): assert [x.scalar.union.stored_type.structure.tag for x in lv.collection.literals] == ["str", "int", "str"] assert v == ["hello", 123, "world"] + +def test_enum_in_dataclass(): + @dataclass_json + @dataclass + class Datum(object): + x: int + y: Color + + lt = TypeEngine.to_literal_type(Datum) + schema = Datum.schema() + schema.fields["y"].load_by = LoadDumpOptions.name + assert lt.metadata == JSONSchema().dump(schema) + + transformer = DataclassTransformer() + ctx = FlyteContext.current_context() + datum = Datum(5, Color.RED) + lv = transformer.to_literal(ctx, datum, Datum, lt) + gt = transformer.guess_python_type(lt) + pv = transformer.to_python_value(ctx, lv, expected_python_type=gt) + assert datum.x == pv.x + assert datum.y.value == pv.y + + @pytest.mark.parametrize( "python_value,python_types,expected_literal_map", [ diff --git a/tests/flytekit/unit/core/test_type_hints.py b/tests/flytekit/unit/core/test_type_hints.py index 1e1f88a756..05a8319a82 100644 --- a/tests/flytekit/unit/core/test_type_hints.py +++ b/tests/flytekit/unit/core/test_type_hints.py @@ -7,6 +7,7 @@ import typing from collections import OrderedDict from dataclasses import dataclass +from enum import Enum from textwrap import dedent import pandas @@ -1067,6 +1068,29 @@ def wf(x: int, y: int) -> Datum: wf(x=10, y=20) +def test_enum_in_dataclass(): + class Color(Enum): + RED = "red" + GREEN = "green" + BLUE = "blue" + + @dataclass_json + @dataclass + class Datum(object): + x: int + y: Color + + @task + def t1(x: int) -> Datum: + return Datum(x=x, y=Color.RED) + + @workflow + def wf(x: int) -> Datum: + return t1(x=x) + + assert wf(x=10) == Datum(10, Color.RED) + + def test_environment(): @task(environment={"FOO": "foofoo", "BAZ": "baz"}) def t1(a: int) -> str: @@ -1463,7 +1487,8 @@ def wf2(a: typing.Union[int, str]) -> typing.Union[int, str]: with pytest.raises( TypeError, - match=dedent(r''' + match=dedent( + r""" Cannot convert from scalar { union { value { @@ -1482,7 +1507,8 @@ def wf2(a: typing.Union[int, str]) -> typing.Union[int, str]: } } to typing.Union\[float, dict\] \(using tag str\) - ''')[1:-1], + """ + )[1:-1], ): assert wf2(a="2") == "2" @@ -1541,7 +1567,9 @@ def __eq__(self, other): "MyInt", MyInt, primitives.Integer.to_flyte_literal_type(), - lambda x: _literal_models.Literal(scalar=_literal_models.Scalar(primitive=_literal_models.Primitive(integer=x.val))), + lambda x: _literal_models.Literal( + scalar=_literal_models.Scalar(primitive=_literal_models.Primitive(integer=x.val)) + ), lambda x: MyInt(x.scalar.primitive.integer), ) ) @@ -1557,8 +1585,7 @@ def wf(a: int) -> int: return t1(a=a) with pytest.raises( - TypeError, - match="Ambiguous choice of variant for union type. Both int and MyInt transformers match" + TypeError, match="Ambiguous choice of variant for union type. Both int and MyInt transformers match" ): assert wf(a=10) == 10 @@ -1580,7 +1607,9 @@ def __eq__(self, other): "MyInt", MyInt, primitives.Integer.to_flyte_literal_type(), - lambda x: _literal_models.Literal(scalar=_literal_models.Scalar(primitive=_literal_models.Primitive(integer=x.val))), + lambda x: _literal_models.Literal( + scalar=_literal_models.Scalar(primitive=_literal_models.Primitive(integer=x.val)) + ), lambda x: MyInt(x.scalar.primitive.integer), ) ) From 9b84040421fc25cc5dbfe371d742c14d0908a2d2 Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Wed, 24 Nov 2021 10:01:25 -0800 Subject: [PATCH 023/128] Fix subworkflow and launch plan FlyteRemote behavior (#751) Signed-off-by: Yee Hing Tong Signed-off-by: maximsmol --- flytekit/clients/friendly.py | 3 +- flytekit/models/core/catalog.py | 75 +++++ flytekit/models/execution.py | 24 +- flytekit/models/node_execution.py | 115 ++++++- flytekit/remote/__init__.py | 7 +- flytekit/remote/component_nodes.py | 21 +- flytekit/remote/executions.py | 239 +++++++++++++++ flytekit/remote/identifier.py | 137 --------- flytekit/remote/interface.py | 13 - flytekit/remote/launch_plan.py | 15 +- flytekit/remote/nodes.py | 177 ++--------- flytekit/remote/remote.py | 283 +++++++++++------- flytekit/remote/{tasks => }/task.py | 3 +- flytekit/remote/tasks/__init__.py | 0 flytekit/remote/tasks/executions.py | 96 ------ flytekit/remote/workflow.py | 99 +++--- flytekit/remote/workflow_execution.py | 76 ----- .../integration/remote/test_remote.py | 4 +- .../unit/models/admin/test_node_executions.py | 44 +++ .../flytekit/unit/models/core/test_catalog.py | 32 ++ tests/flytekit/unit/remote/test_identifier.py | 77 ----- tests/flytekit/unit/remote/test_remote.py | 96 +----- .../unit/remote/test_wrapper_classes.py | 5 - 23 files changed, 811 insertions(+), 830 deletions(-) create mode 100644 flytekit/models/core/catalog.py create mode 100644 flytekit/remote/executions.py delete mode 100644 flytekit/remote/identifier.py rename flytekit/remote/{tasks => }/task.py (95%) delete mode 100644 flytekit/remote/tasks/__init__.py delete mode 100644 flytekit/remote/tasks/executions.py create mode 100644 tests/flytekit/unit/models/core/test_catalog.py delete mode 100644 tests/flytekit/unit/remote/test_identifier.py diff --git a/flytekit/clients/friendly.py b/flytekit/clients/friendly.py index a34bf31a48..5a6bbc4fdc 100644 --- a/flytekit/clients/friendly.py +++ b/flytekit/clients/friendly.py @@ -676,12 +676,11 @@ def get_node_execution(self, node_execution_identifier): ) ) - def get_node_execution_data(self, node_execution_identifier): + def get_node_execution_data(self, node_execution_identifier) -> _execution.NodeExecutionGetDataResponse: """ Returns signed URLs to LiteralMap blobs for a node execution's inputs and outputs (when available). :param flytekit.models.core.identifier.NodeExecutionIdentifier node_execution_identifier: - :rtype: flytekit.models.execution.NodeExecutionGetDataResponse """ return _execution.NodeExecutionGetDataResponse.from_flyte_idl( super(SynchronousFlyteClient, self).get_node_execution_data( diff --git a/flytekit/models/core/catalog.py b/flytekit/models/core/catalog.py new file mode 100644 index 0000000000..97cc2c34ff --- /dev/null +++ b/flytekit/models/core/catalog.py @@ -0,0 +1,75 @@ +from flyteidl.core import catalog_pb2 + +from flytekit.models import common as _common_models +from flytekit.models.core import identifier as _identifier + + +class CatalogArtifactTag(_common_models.FlyteIdlEntity): + def __init__(self, artifact_id: str, name: str): + self._artifact_id = artifact_id + self._name = name + + @property + def artifact_id(self) -> str: + return self._artifact_id + + @property + def name(self) -> str: + return self._name + + def to_flyte_idl(self) -> catalog_pb2.CatalogArtifactTag: + return catalog_pb2.CatalogArtifactTag(artifact_id=self.artifact_id, name=self.name) + + @classmethod + def from_flyte_idl(cls, p: catalog_pb2.CatalogArtifactTag) -> "CatalogArtifactTag": + return cls( + artifact_id=p.artifact_id, + name=p.name, + ) + + +class CatalogMetadata(_common_models.FlyteIdlEntity): + def __init__( + self, + dataset_id: _identifier.Identifier, + artifact_tag: CatalogArtifactTag, + source_task_execution: _identifier.TaskExecutionIdentifier, + ): + self._dataset_id = dataset_id + self._artifact_tag = artifact_tag + self._source_task_execution = source_task_execution + + @property + def dataset_id(self) -> _identifier.Identifier: + return self._dataset_id + + @property + def artifact_tag(self) -> CatalogArtifactTag: + return self._artifact_tag + + @property + def source_task_execution(self) -> _identifier.TaskExecutionIdentifier: + return self._source_task_execution + + @property + def source_execution(self) -> _identifier.TaskExecutionIdentifier: + """ + This is a one of but for now there's only one thing in the one of + """ + return self._source_task_execution + + def to_flyte_idl(self) -> catalog_pb2.CatalogMetadata: + return catalog_pb2.CatalogMetadata( + dataset_id=self.dataset_id.to_flyte_idl(), + artifact_tag=self.artifact_tag.to_flyte_idl(), + source_task_execution=self.source_task_execution.to_flyte_idl(), + ) + + @classmethod + def from_flyte_idl(cls, pb: catalog_pb2.CatalogMetadata) -> "CatalogMetadata": + return cls( + dataset_id=_identifier.Identifier.from_flyte_idl(pb.dataset_id), + artifact_tag=CatalogArtifactTag.from_flyte_idl(pb.artifact_tag), + # Add HasField check if more things are ever added to the one of + source_task_execution=_identifier.TaskExecutionIdentifier.from_flyte_idl(pb.source_task_execution), + ) diff --git a/flytekit/models/execution.py b/flytekit/models/execution.py index 042b3ad909..a29bad80f1 100644 --- a/flytekit/models/execution.py +++ b/flytekit/models/execution.py @@ -1,3 +1,5 @@ +import typing + import flyteidl.admin.execution_pb2 as _execution_pb2 import flyteidl.admin.node_execution_pb2 as _node_execution_pb2 import flyteidl.admin.task_execution_pb2 as _task_execution_pb2 @@ -7,6 +9,7 @@ from flytekit.models import literals as _literals_models from flytekit.models.core import execution as _core_execution from flytekit.models.core import identifier as _identifier +from flytekit.models.node_execution import DynamicWorkflowNodeMetadata class ExecutionMetadata(_common_models.FlyteIdlEntity): @@ -238,7 +241,6 @@ class Execution(_common_models.FlyteIdlEntity): def __init__(self, id, spec, closure): """ :param flytekit.models.core.identifier.WorkflowExecutionIdentifier id: - :param Text id: :param ExecutionSpec spec: :param ExecutionClosure closure: """ @@ -403,8 +405,8 @@ def __init__(self, inputs, outputs, full_inputs, full_outputs): """ :param _common_models.UrlBlob inputs: :param _common_models.UrlBlob outputs: - :param _literals_pb2.LiteralMap full_inputs: - :param _literals_pb2.LiteralMap full_outputs: + :param _literals_models.LiteralMap full_inputs: + :param _literals_models.LiteralMap full_outputs: """ self._inputs = inputs self._outputs = outputs @@ -428,14 +430,14 @@ def outputs(self): @property def full_inputs(self): """ - :rtype: _literals_pb2.LiteralMap + :rtype: _literals_models.LiteralMap """ return self._full_inputs @property def full_outputs(self): """ - :rtype: _literals_pb2.LiteralMap + :rtype: _literals_models.LiteralMap """ return self._full_outputs @@ -493,6 +495,14 @@ def to_flyte_idl(self): class NodeExecutionGetDataResponse(_CommonDataResponse): + def __init__(self, *args, dynamic_workflow: typing.Optional[DynamicWorkflowNodeMetadata] = None, **kwargs): + super().__init__(*args, **kwargs) + self._dynamic_workflow = dynamic_workflow + + @property + def dynamic_workflow(self) -> typing.Optional[DynamicWorkflowNodeMetadata]: + return self._dynamic_workflow + @classmethod def from_flyte_idl(cls, pb2_object): """ @@ -504,6 +514,9 @@ def from_flyte_idl(cls, pb2_object): outputs=_common_models.UrlBlob.from_flyte_idl(pb2_object.outputs), full_inputs=_literals_models.LiteralMap.from_flyte_idl(pb2_object.full_inputs), full_outputs=_literals_models.LiteralMap.from_flyte_idl(pb2_object.full_outputs), + dynamic_workflow=DynamicWorkflowNodeMetadata.from_flyte_idl(pb2_object.dynamic_workflow) + if pb2_object.HasField("dynamic_workflow") + else None, ) def to_flyte_idl(self): @@ -515,4 +528,5 @@ def to_flyte_idl(self): outputs=self.outputs.to_flyte_idl(), full_inputs=self.full_inputs.to_flyte_idl(), full_outputs=self.full_outputs.to_flyte_idl(), + dynamic_workflow=self.dynamic_workflow.to_flyte_idl() if self.dynamic_workflow else None, ) diff --git a/flytekit/models/node_execution.py b/flytekit/models/node_execution.py index 5a0cda6a6e..762dfd196f 100644 --- a/flytekit/models/node_execution.py +++ b/flytekit/models/node_execution.py @@ -1,13 +1,101 @@ +import typing + import flyteidl.admin.node_execution_pb2 as _node_execution_pb2 import pytz as _pytz from flytekit.models import common as _common_models +from flytekit.models.core import catalog as catalog_models +from flytekit.models.core import compiler as core_compiler_models from flytekit.models.core import execution as _core_execution from flytekit.models.core import identifier as _identifier +class WorkflowNodeMetadata(_common_models.FlyteIdlEntity): + def __init__(self, execution_id: _identifier.WorkflowExecutionIdentifier): + self._execution_id = execution_id + + @property + def execution_id(self) -> _identifier.WorkflowExecutionIdentifier: + return self._execution_id + + def to_flyte_idl(self) -> _node_execution_pb2.WorkflowNodeMetadata: + return _node_execution_pb2.WorkflowNodeMetadata( + executionId=self.execution_id.to_flyte_idl(), + ) + + @classmethod + def from_flyte_idl(cls, p: _node_execution_pb2.WorkflowNodeMetadata) -> "WorkflowNodeMetadata": + return cls( + execution_id=_identifier.WorkflowExecutionIdentifier.from_flyte_idl(p.executionId), + ) + + +class DynamicWorkflowNodeMetadata(_common_models.FlyteIdlEntity): + def __init__(self, id: _identifier.Identifier, compiled_workflow: core_compiler_models.CompiledWorkflowClosure): + self._id = id + self._compiled_workflow = compiled_workflow + + @property + def id(self) -> _identifier.Identifier: + return self._id + + @property + def compiled_workflow(self) -> core_compiler_models.CompiledWorkflowClosure: + return self._compiled_workflow + + def to_flyte_idl(self) -> _node_execution_pb2.DynamicWorkflowNodeMetadata: + return _node_execution_pb2.DynamicWorkflowNodeMetadata( + id=self.id.to_flyte_idl(), + compiled_workflow=self.compiled_workflow.to_flyte_idl(), + ) + + @classmethod + def from_flyte_idl(cls, p: _node_execution_pb2.DynamicWorkflowNodeMetadata) -> "DynamicWorkflowNodeMetadata": + yy = cls( + id=_identifier.Identifier.from_flyte_idl(p.id), + compiled_workflow=core_compiler_models.CompiledWorkflowClosure.from_flyte_idl(p.compiled_workflow), + ) + return yy + + +class TaskNodeMetadata(_common_models.FlyteIdlEntity): + def __init__(self, cache_status: int, catalog_key: catalog_models.CatalogMetadata): + self._cache_status = cache_status + self._catalog_key = catalog_key + + @property + def cache_status(self) -> int: + return self._cache_status + + @property + def catalog_key(self) -> catalog_models.CatalogMetadata: + return self._catalog_key + + def to_flyte_idl(self) -> _node_execution_pb2.TaskNodeMetadata: + return _node_execution_pb2.TaskNodeMetadata( + cache_status=self.cache_status, + catalog_key=self.catalog_key.to_flyte_idl(), + ) + + @classmethod + def from_flyte_idl(cls, p: _node_execution_pb2.TaskNodeMetadata) -> "TaskNodeMetadata": + return cls( + cache_status=p.cache_status, + catalog_key=catalog_models.CatalogMetadata.from_flyte_idl(p.catalog_key), + ) + + class NodeExecutionClosure(_common_models.FlyteIdlEntity): - def __init__(self, phase, started_at, duration, output_uri=None, error=None): + def __init__( + self, + phase, + started_at, + duration, + output_uri=None, + error=None, + workflow_node_metadata: typing.Optional[WorkflowNodeMetadata] = None, + task_node_metadata: typing.Optional[TaskNodeMetadata] = None, + ): """ :param int phase: :param datetime.datetime started_at: @@ -20,6 +108,9 @@ def __init__(self, phase, started_at, duration, output_uri=None, error=None): self._duration = duration self._output_uri = output_uri self._error = error + self._workflow_node_metadata = workflow_node_metadata + self._task_node_metadata = task_node_metadata + # TODO: Add output_data field as well. @property def phase(self): @@ -56,6 +147,18 @@ def error(self): """ return self._error + @property + def workflow_node_metadata(self) -> typing.Optional[WorkflowNodeMetadata]: + return self._workflow_node_metadata + + @property + def task_node_metadata(self) -> typing.Optional[TaskNodeMetadata]: + return self._task_node_metadata + + @property + def target_metadata(self) -> typing.Union[WorkflowNodeMetadata, TaskNodeMetadata]: + return self.workflow_node_metadata or self.task_node_metadata + def to_flyte_idl(self): """ :rtype: flyteidl.admin.node_execution_pb2.NodeExecutionClosure @@ -64,6 +167,10 @@ def to_flyte_idl(self): phase=self.phase, output_uri=self.output_uri, error=self.error.to_flyte_idl() if self.error is not None else None, + workflow_node_metadata=self.workflow_node_metadata.to_flyte_idl() + if self.workflow_node_metadata is not None + else None, + task_node_metadata=self.task_node_metadata.to_flyte_idl() if self.task_node_metadata is not None else None, ) obj.started_at.FromDatetime(self.started_at.astimezone(_pytz.UTC).replace(tzinfo=None)) obj.duration.FromTimedelta(self.duration) @@ -81,6 +188,12 @@ def from_flyte_idl(cls, p): error=_core_execution.ExecutionError.from_flyte_idl(p.error) if p.HasField("error") else None, started_at=p.started_at.ToDatetime().replace(tzinfo=_pytz.UTC), duration=p.duration.ToTimedelta(), + workflow_node_metadata=WorkflowNodeMetadata.from_flyte_idl(p.workflow_node_metadata) + if p.HasField("workflow_node_metadata") + else None, + task_node_metadata=TaskNodeMetadata.from_flyte_idl(p.task_node_metadata) + if p.HasField("task_node_metadata") + else None, ) diff --git a/flytekit/remote/__init__.py b/flytekit/remote/__init__.py index 9dc4a5f0ed..51def80359 100644 --- a/flytekit/remote/__init__.py +++ b/flytekit/remote/__init__.py @@ -79,10 +79,9 @@ """ from flytekit.remote.component_nodes import FlyteTaskNode, FlyteWorkflowNode +from flytekit.remote.executions import FlyteNodeExecution, FlyteTaskExecution, FlyteWorkflowExecution from flytekit.remote.launch_plan import FlyteLaunchPlan -from flytekit.remote.nodes import FlyteNode, FlyteNodeExecution +from flytekit.remote.nodes import FlyteNode from flytekit.remote.remote import FlyteRemote -from flytekit.remote.tasks.executions import FlyteTaskExecution -from flytekit.remote.tasks.task import FlyteTask +from flytekit.remote.task import FlyteTask from flytekit.remote.workflow import FlyteWorkflow -from flytekit.remote.workflow_execution import FlyteWorkflowExecution diff --git a/flytekit/remote/component_nodes.py b/flytekit/remote/component_nodes.py index 06b885abfd..367cab8997 100644 --- a/flytekit/remote/component_nodes.py +++ b/flytekit/remote/component_nodes.py @@ -1,23 +1,22 @@ import logging as _logging from typing import Dict -import flytekit from flytekit.common.exceptions import system as _system_exceptions from flytekit.models import launch_plan as _launch_plan_model from flytekit.models import task as _task_model +from flytekit.models.core import identifier as id_models from flytekit.models.core import workflow as _workflow_model -from flytekit.remote import identifier as _identifier class FlyteTaskNode(_workflow_model.TaskNode): """A class encapsulating a task that a Flyte node needs to execute.""" - def __init__(self, flyte_task: "flytekit.remote.tasks.task.FlyteTask"): + def __init__(self, flyte_task: "flytekit.remote.task.FlyteTask"): self._flyte_task = flyte_task super(FlyteTaskNode, self).__init__(None) @property - def reference_id(self) -> _identifier.Identifier: + def reference_id(self) -> id_models.Identifier: """A globally unique identifier for the task.""" return self._flyte_task.id @@ -29,7 +28,7 @@ def flyte_task(self) -> "flytekit.remote.tasks.task.FlyteTask": def promote_from_model( cls, base_model: _workflow_model.TaskNode, - tasks: Dict[_identifier.Identifier, _task_model.TaskTemplate], + tasks: Dict[id_models.Identifier, _task_model.TaskTemplate], ) -> "FlyteTaskNode": """ Takes the idl wrapper for a TaskNode and returns the hydrated Flytekit object for it by fetching it with the @@ -38,12 +37,12 @@ def promote_from_model( :param base_model: :param tasks: """ - from flytekit.remote.tasks import task as _task + from flytekit.remote.task import FlyteTask if base_model.reference_id in tasks: task = tasks[base_model.reference_id] _logging.info(f"Found existing task template for {task.id}, will not retrieve from Admin") - flyte_task = _task.FlyteTask.promote_from_model(task) + flyte_task = FlyteTask.promote_from_model(task) return cls(flyte_task) raise _system_exceptions.FlyteSystemException(f"Task template {base_model.reference_id} not found.") @@ -76,7 +75,7 @@ def __repr__(self) -> str: return f"FlyteWorkflowNode with launch plan: {self.flyte_launch_plan}" @property - def launchplan_ref(self) -> _identifier.Identifier: + def launchplan_ref(self) -> id_models.Identifier: """A globally unique identifier for the launch plan, which should map to Admin.""" return self._flyte_launch_plan.id if self._flyte_launch_plan else None @@ -96,9 +95,9 @@ def flyte_workflow(self) -> "flytekit.remote.workflow.FlyteWorkflow": def promote_from_model( cls, base_model: _workflow_model.WorkflowNode, - sub_workflows: Dict[_identifier.Identifier, _workflow_model.WorkflowTemplate], - node_launch_plans: Dict[_identifier.Identifier, _launch_plan_model.LaunchPlanSpec], - tasks: Dict[_identifier.Identifier, _task_model.TaskTemplate], + sub_workflows: Dict[id_models.Identifier, _workflow_model.WorkflowTemplate], + node_launch_plans: Dict[id_models.Identifier, _launch_plan_model.LaunchPlanSpec], + tasks: Dict[id_models.Identifier, _task_model.TaskTemplate], ) -> "FlyteWorkflowNode": from flytekit.remote import launch_plan as _launch_plan from flytekit.remote import workflow as _workflow diff --git a/flytekit/remote/executions.py b/flytekit/remote/executions.py new file mode 100644 index 0000000000..05b94b2302 --- /dev/null +++ b/flytekit/remote/executions.py @@ -0,0 +1,239 @@ +from __future__ import annotations + +from typing import Any, Dict, List, Optional, Union + +from flytekit.common.exceptions import user as _user_exceptions +from flytekit.common.exceptions import user as user_exceptions +from flytekit.models import execution as execution_models +from flytekit.models import node_execution as node_execution_models +from flytekit.models.admin import task_execution as admin_task_execution_models +from flytekit.models.core import execution as core_execution_models +from flytekit.remote.workflow import FlyteWorkflow + + +class FlyteTaskExecution(admin_task_execution_models.TaskExecution): + """A class encapsulating a task execution being run on a Flyte remote backend.""" + + def __init__(self, *args, **kwargs): + super(FlyteTaskExecution, self).__init__(*args, **kwargs) + self._inputs = None + self._outputs = None + + @property + def is_complete(self) -> bool: + """Whether or not the execution is complete.""" + return self.closure.phase in { + core_execution_models.TaskExecutionPhase.ABORTED, + core_execution_models.TaskExecutionPhase.FAILED, + core_execution_models.TaskExecutionPhase.SUCCEEDED, + } + + @property + def inputs(self) -> Dict[str, Any]: + """ + Returns the inputs of the task execution in the standard Python format that is produced by + the type engine. + """ + return self._inputs + + @property + def outputs(self) -> Dict[str, Any]: + """ + Returns the outputs of the task execution, if available, in the standard Python format that is produced by + the type engine. + + :raises: ``FlyteAssertion`` error if execution is in progress or execution ended in error. + """ + if not self.is_complete: + raise user_exceptions.FlyteAssertion( + "Please wait until the node execution has completed before requesting the outputs." + ) + if self.error: + raise user_exceptions.FlyteAssertion("Outputs could not be found because the execution ended in failure.") + return self._outputs + + @property + def error(self) -> Optional[core_execution_models.ExecutionError]: + """ + If execution is in progress, raise an exception. Otherwise, return None if no error was present upon + reaching completion. + """ + if not self.is_complete: + raise user_exceptions.FlyteAssertion( + "Please what until the task execution has completed before requesting error information." + ) + return self.closure.error + + @classmethod + def promote_from_model(cls, base_model: admin_task_execution_models.TaskExecution) -> "FlyteTaskExecution": + return cls( + closure=base_model.closure, + id=base_model.id, + input_uri=base_model.input_uri, + is_parent=base_model.is_parent, + ) + + +class FlyteWorkflowExecution(execution_models.Execution): + """A class encapsulating a workflow execution being run on a Flyte remote backend.""" + + def __init__(self, *args, **kwargs): + super(FlyteWorkflowExecution, self).__init__(*args, **kwargs) + self._node_executions = None + self._inputs = None + self._outputs = None + self._flyte_workflow: Optional[FlyteWorkflow] = None + + @property + def node_executions(self) -> Dict[str, "FlyteNodeExecution"]: + """Get a dictionary of node executions that are a part of this workflow execution.""" + return self._node_executions or {} + + @property + def inputs(self) -> Dict[str, Any]: + """ + Returns the inputs to the execution in the standard python format as dictated by the type engine. + """ + return self._inputs + + @property + def outputs(self) -> Dict[str, Any]: + """ + Returns the outputs to the execution in the standard python format as dictated by the type engine. + + :raises: ``FlyteAssertion`` error if execution is in progress or execution ended in error. + """ + if not self.is_complete: + raise _user_exceptions.FlyteAssertion( + "Please wait until the node execution has completed before requesting the outputs." + ) + if self.error: + raise _user_exceptions.FlyteAssertion("Outputs could not be found because the execution ended in failure.") + return self._outputs + + @property + def error(self) -> core_execution_models.ExecutionError: + """ + If execution is in progress, raise an exception. Otherwise, return None if no error was present upon + reaching completion. + """ + if not self.is_complete: + raise _user_exceptions.FlyteAssertion( + "Please wait until a workflow has completed before checking for an error." + ) + return self.closure.error + + @property + def is_complete(self) -> bool: + """ + Whether or not the execution is complete. + """ + return self.closure.phase in { + core_execution_models.WorkflowExecutionPhase.ABORTED, + core_execution_models.WorkflowExecutionPhase.FAILED, + core_execution_models.WorkflowExecutionPhase.SUCCEEDED, + core_execution_models.WorkflowExecutionPhase.TIMED_OUT, + } + + @classmethod + def promote_from_model(cls, base_model: execution_models.Execution) -> "FlyteWorkflowExecution": + return cls( + closure=base_model.closure, + id=base_model.id, + spec=base_model.spec, + ) + + +class FlyteNodeExecution(node_execution_models.NodeExecution): + """A class encapsulating a node execution being run on a Flyte remote backend.""" + + def __init__(self, *args, **kwargs): + super(FlyteNodeExecution, self).__init__(*args, **kwargs) + self._task_executions = None + self._workflow_executions = [] + self._underlying_node_executions = None + self._inputs = None + self._outputs = None + self._interface = None + + @property + def task_executions(self) -> List[FlyteTaskExecution]: + return self._task_executions or [] + + @property + def workflow_executions(self) -> List[FlyteWorkflowExecution]: + return self._workflow_executions + + @property + def subworkflow_node_executions(self) -> Dict[str, FlyteNodeExecution]: + """ + This returns underlying node executions in instances where the current node execution is + a parent node. This happens when it's either a static or dynamic subworkflow. + """ + return ( + {} + if self._underlying_node_executions is None + else {n.id.node_id: n for n in self._underlying_node_executions} + ) + + @property + def executions(self) -> List[Union[FlyteTaskExecution, FlyteWorkflowExecution]]: + return self.task_executions or self._underlying_node_executions or [] + + @property + def inputs(self) -> Dict[str, Any]: + """ + Returns the inputs to the execution in the standard python format as dictated by the type engine. + """ + return self._inputs + + @property + def outputs(self) -> Dict[str, Any]: + """ + Returns the outputs to the execution in the standard python format as dictated by the type engine. + + :raises: ``FlyteAssertion`` error if execution is in progress or execution ended in error. + """ + if not self.is_complete: + raise _user_exceptions.FlyteAssertion( + "Please wait until the node execution has completed before requesting the outputs." + ) + if self.error: + raise _user_exceptions.FlyteAssertion("Outputs could not be found because the execution ended in failure.") + return self._outputs + + @property + def error(self) -> core_execution_models.ExecutionError: + """ + If execution is in progress, raise an exception. Otherwise, return None if no error was present upon + reaching completion. + """ + if not self.is_complete: + raise _user_exceptions.FlyteAssertion( + "Please wait until the node execution has completed before requesting error information." + ) + return self.closure.error + + @property + def is_complete(self) -> bool: + """Whether or not the execution is complete.""" + return self.closure.phase in { + core_execution_models.NodeExecutionPhase.ABORTED, + core_execution_models.NodeExecutionPhase.FAILED, + core_execution_models.NodeExecutionPhase.SKIPPED, + core_execution_models.NodeExecutionPhase.SUCCEEDED, + core_execution_models.NodeExecutionPhase.TIMED_OUT, + } + + @classmethod + def promote_from_model(cls, base_model: node_execution_models.NodeExecution) -> "FlyteNodeExecution": + return cls( + closure=base_model.closure, id=base_model.id, input_uri=base_model.input_uri, metadata=base_model.metadata + ) + + @property + def interface(self) -> "flytekit.remote.interface.TypedInterface": + """ + Return the interface of the task or subworkflow associated with this node execution. + """ + return self._interface diff --git a/flytekit/remote/identifier.py b/flytekit/remote/identifier.py deleted file mode 100644 index 611c9af639..0000000000 --- a/flytekit/remote/identifier.py +++ /dev/null @@ -1,137 +0,0 @@ -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.models.core import identifier as _core_identifier - - -class Identifier(_core_identifier.Identifier): - - _STRING_TO_TYPE_MAP = { - "lp": _core_identifier.ResourceType.LAUNCH_PLAN, - "wf": _core_identifier.ResourceType.WORKFLOW, - "tsk": _core_identifier.ResourceType.TASK, - } - _TYPE_TO_STRING_MAP = {v: k for k, v in _STRING_TO_TYPE_MAP.items()} - - @classmethod - def promote_from_model(cls, base_model: _core_identifier.Identifier) -> "Identifier": - return cls(base_model.resource_type, base_model.project, base_model.domain, base_model.name, base_model.version) - - @classmethod - def from_urn(cls, urn: str) -> "Identifier": - """ - Parses a string urn in the correct format into an identifier - """ - segments = urn.split(":") - if len(segments) != 5: - raise _user_exceptions.FlyteValueException( - urn, - "The provided string was not in a parseable format. The string for an identifier must be in the " - "format entity_type:project:domain:name:version.", - ) - - resource_type, project, domain, name, version = segments - - if resource_type not in cls._STRING_TO_TYPE_MAP: - raise _user_exceptions.FlyteValueException( - resource_type, - "The provided string could not be parsed. The first element of an identifier must be one of: " - f"{list(cls._STRING_TO_TYPE_MAP.keys())}. ", - ) - - return cls(cls._STRING_TO_TYPE_MAP[resource_type], project, domain, name, version) - - def __str__(self): - return ( - f"{type(self)._TYPE_TO_STRING_MAP.get(self.resource_type, '')}:" - f"{self.project}:" - f"{self.domain}:" - f"{self.name}:" - f"{self.version}" - ) - - -class WorkflowExecutionIdentifier(_core_identifier.WorkflowExecutionIdentifier): - @classmethod - def promote_from_model( - cls, base_model: _core_identifier.WorkflowExecutionIdentifier - ) -> "WorkflowExecutionIdentifier": - return cls(base_model.project, base_model.domain, base_model.name) - - @classmethod - def from_urn(cls, string: str) -> "WorkflowExecutionIdentifier": - """ - Parses a string in the correct format into an identifier - """ - segments = string.split(":") - if len(segments) != 4: - raise _user_exceptions.FlyteValueException( - string, - "The provided string was not in a parseable format. The string for an identifier must be in the format" - " ex:project:domain:name.", - ) - - resource_type, project, domain, name = segments - - if resource_type != "ex": - raise _user_exceptions.FlyteValueException( - resource_type, - "The provided string could not be parsed. The first element of an execution identifier must be 'ex'.", - ) - - return cls(project, domain, name) - - def __str__(self): - return f"ex:{self.project}:{self.domain}:{self.name}" - - -class TaskExecutionIdentifier(_core_identifier.TaskExecutionIdentifier): - @classmethod - def promote_from_model(cls, base_model: _core_identifier.TaskExecutionIdentifier) -> "TaskExecutionIdentifier": - return cls( - task_id=base_model.task_id, - node_execution_id=base_model.node_execution_id, - retry_attempt=base_model.retry_attempt, - ) - - @classmethod - def from_urn(cls, string: str) -> "TaskExecutionIdentifier": - """ - Parses a string in the correct format into an identifier - """ - segments = string.split(":") - if len(segments) != 10: - raise _user_exceptions.FlyteValueException( - string, - "The provided string was not in a parseable format. The string for an identifier must be in the format" - " te:exec_project:exec_domain:exec_name:node_id:task_project:task_domain:task_name:task_version:retry.", - ) - - resource_type, ep, ed, en, node_id, tp, td, tn, tv, retry = segments - - if resource_type != "te": - raise _user_exceptions.FlyteValueException( - resource_type, - "The provided string could not be parsed. The first element of an execution identifier must be 'ex'.", - ) - - return cls( - task_id=Identifier(_core_identifier.ResourceType.TASK, tp, td, tn, tv), - node_execution_id=_core_identifier.NodeExecutionIdentifier( - node_id=node_id, - execution_id=_core_identifier.WorkflowExecutionIdentifier(ep, ed, en), - ), - retry_attempt=int(retry), - ) - - def __str__(self): - return ( - "te:" - f"{self.node_execution_id.execution_id.project}:" - f"{self.node_execution_id.execution_id.domain}:" - f"{self.node_execution_id.execution_id.name}:" - f"{self.node_execution_id.node_id}:" - f"{self.task_id.project}:" - f"{self.task_id.domain}:" - f"{self.task_id.name}:" - f"{self.task_id.version}:" - f"{self.retry_attempt}" - ) diff --git a/flytekit/remote/interface.py b/flytekit/remote/interface.py index 6aeb0e236d..df61c8e336 100644 --- a/flytekit/remote/interface.py +++ b/flytekit/remote/interface.py @@ -1,8 +1,4 @@ -from typing import Any, Dict, List, Tuple - from flytekit.models import interface as _interface_models -from flytekit.models import literals as _literal_models -from flytekit.remote import nodes as _nodes class TypedInterface(_interface_models.TypedInterface): @@ -13,12 +9,3 @@ def promote_from_model(cls, model): :rtype: TypedInterface """ return cls(model.inputs, model.outputs) - - def create_bindings_for_inputs( - self, map_of_bindings: Dict[str, Any] - ) -> Tuple[List[_literal_models.Binding], List[_nodes.FlyteNode]]: - """ - :param: map_of_bindings: this can be scalar primitives, it can be node output references, lists, etc. - :raises: flytekit.common.exceptions.user.FlyteAssertion - """ - return [], [] diff --git a/flytekit/remote/launch_plan.py b/flytekit/remote/launch_plan.py index 200244f394..016e3a3489 100644 --- a/flytekit/remote/launch_plan.py +++ b/flytekit/remote/launch_plan.py @@ -7,8 +7,7 @@ from flytekit.engines.flyte import engine as _flyte_engine from flytekit.models import interface as _interface_models from flytekit.models import launch_plan as _launch_plan_models -from flytekit.models.core import identifier as _identifier_model -from flytekit.remote import identifier as _identifier +from flytekit.models.core import identifier as id_models from flytekit.remote import interface as _interface @@ -27,11 +26,11 @@ def __init__(self, id, *args, **kwargs): @classmethod def promote_from_model( - cls, id: _identifier.Identifier, model: _launch_plan_models.LaunchPlanSpec + cls, id: id_models.Identifier, model: _launch_plan_models.LaunchPlanSpec ) -> "FlyteLaunchPlan": lp = cls( id=id, - workflow_id=_identifier.Identifier.promote_from_model(model.workflow_id), + workflow_id=model.workflow_id, default_inputs=_interface_models.ParameterMap(model.default_inputs.parameters), fixed_inputs=model.fixed_inputs, entity_metadata=model.entity_metadata, @@ -50,7 +49,7 @@ def promote_from_model( return lp @property - def id(self) -> _identifier.Identifier: + def id(self) -> id_models.Identifier: return self._id @property @@ -65,7 +64,7 @@ def is_scheduled(self) -> bool: return False @property - def workflow_id(self) -> _identifier.Identifier: + def workflow_id(self) -> id_models.Identifier: return self._workflow_id @property @@ -78,8 +77,8 @@ def interface(self) -> _interface.TypedInterface: return self._interface @property - def resource_type(self) -> _identifier_model.ResourceType: - return _identifier_model.ResourceType.LAUNCH_PLAN + def resource_type(self) -> id_models.ResourceType: + return id_models.ResourceType.LAUNCH_PLAN @property def entity_type_text(self) -> str: diff --git a/flytekit/remote/nodes.py b/flytekit/remote/nodes.py index 68d84f00ea..f8ae1b2d6a 100644 --- a/flytekit/remote/nodes.py +++ b/flytekit/remote/nodes.py @@ -1,24 +1,18 @@ +from __future__ import annotations + import logging as _logging -from typing import Any, Dict, List, Optional, Union +from typing import Dict, List, Optional, Union -import flytekit -from flytekit.clients.helpers import iterate_node_executions, iterate_task_executions from flytekit.common import constants as _constants from flytekit.common.exceptions import system as _system_exceptions from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.mixins import artifact as _artifact_mixin from flytekit.common.mixins import hash as _hash_mixin -from flytekit.common.utils import _dnsify from flytekit.core.promise import NodeOutput -from flytekit.engines.flyte import engine as _flyte_engine from flytekit.models import launch_plan as _launch_plan_model -from flytekit.models import node_execution as _node_execution_models from flytekit.models import task as _task_model -from flytekit.models.core import execution as _execution_models +from flytekit.models.core import identifier as id_models from flytekit.models.core import workflow as _workflow_model from flytekit.remote import component_nodes as _component_nodes -from flytekit.remote import identifier as _identifier -from flytekit.remote.tasks.executions import FlyteTaskExecution class FlyteNode(_hash_mixin.HashOnReferenceMixin, _workflow_model.Node): @@ -34,7 +28,6 @@ def __init__( flyte_workflow: Optional["FlyteWorkflow"] = None, flyte_launch_plan: Optional["FlyteLaunchPlan"] = None, flyte_branch=None, - parameter_mapping=True, ): non_none_entities = list(filter(None, [flyte_task, flyte_workflow, flyte_launch_plan, flyte_branch])) if len(non_none_entities) != 1: @@ -50,15 +43,20 @@ def __init__( elif flyte_launch_plan is not None: workflow_node = _component_nodes.FlyteWorkflowNode(flyte_launch_plan=flyte_launch_plan) + task_node = None + if flyte_task: + task_node = _component_nodes.FlyteTaskNode(flyte_task) + branch_node = None + super(FlyteNode, self).__init__( - id=_dnsify(id) if id else None, + id=id, metadata=metadata, inputs=bindings, upstream_node_ids=[n.id for n in upstream_nodes], output_aliases=[], - task_node=_component_nodes.FlyteTaskNode(flyte_task) if flyte_task else None, + task_node=task_node, workflow_node=workflow_node, - branch_node=flyte_branch, + branch_node=branch_node, ) self._upstream = upstream_nodes @@ -70,11 +68,12 @@ def flyte_entity(self) -> Union["FlyteTask", "FlyteWorkflow", "FlyteLaunchPlan"] def promote_from_model( cls, model: _workflow_model.Node, - sub_workflows: Optional[Dict[_identifier.Identifier, _workflow_model.WorkflowTemplate]], - node_launch_plans: Optional[Dict[_identifier.Identifier, _launch_plan_model.LaunchPlanSpec]], - tasks: Optional[Dict[_identifier.Identifier, _task_model.TaskTemplate]], - ) -> "FlyteNode": - id = model.id + sub_workflows: Optional[Dict[id_models.Identifier, _workflow_model.WorkflowTemplate]], + node_launch_plans: Optional[Dict[id_models.Identifier, _launch_plan_model.LaunchPlanSpec]], + tasks: Optional[Dict[id_models.Identifier, _task_model.TaskTemplate]], + ) -> FlyteNode: + node_model_id = model.id + # TODO: Consider removing if id in {_constants.START_NODE_ID, _constants.END_NODE_ID}: _logging.warning(f"Should not call promote from model on a start node or end node {model}") return None @@ -97,6 +96,7 @@ def promote_from_model( # When WorkflowTemplate models (containing node models) are returned by Admin, they've been compiled with a # start node. In order to make the promoted FlyteWorkflow look the same, we strip the start-node text back out. + # TODO: Consider removing for model_input in model.inputs: if ( model_input.binding.promise is not None @@ -106,7 +106,7 @@ def promote_from_model( if flyte_task_node is not None: return cls( - id=id, + id=node_model_id, upstream_nodes=[], # set downstream, model doesn't contain this information bindings=model.inputs, metadata=model.metadata, @@ -115,7 +115,7 @@ def promote_from_model( elif flyte_workflow_node is not None: if flyte_workflow_node.flyte_workflow is not None: return cls( - id=id, + id=node_model_id, upstream_nodes=[], # set downstream, model doesn't contain this information bindings=model.inputs, metadata=model.metadata, @@ -123,7 +123,7 @@ def promote_from_model( ) elif flyte_workflow_node.flyte_launch_plan is not None: return cls( - id=id, + id=node_model_id, upstream_nodes=[], # set downstream, model doesn't contain this information bindings=model.inputs, metadata=model.metadata, @@ -135,7 +135,7 @@ def promote_from_model( raise _system_exceptions.FlyteSystemException("Bad FlyteNode model, both task and workflow nodes are empty") @property - def upstream_nodes(self) -> List["FlyteNode"]: + def upstream_nodes(self) -> List[FlyteNode]: return self._upstream @property @@ -146,136 +146,5 @@ def upstream_node_ids(self) -> List[str]: def outputs(self) -> Dict[str, NodeOutput]: return self._outputs - def assign_id_and_return(self, id: str): - if self.id: - raise _user_exceptions.FlyteAssertion( - f"Error assigning ID: {id} because {self} is already assigned. Has this node been ssigned to another " - "workflow already?" - ) - self._id = _dnsify(id) if id else None - self._metadata.name = id - return self - - def with_overrides(self, *args, **kwargs): - # TODO: Implement overrides - raise NotImplementedError("Overrides are not supported in Flyte yet.") - def __repr__(self) -> str: return f"Node(ID: {self.id})" - - -class FlyteNodeExecution(_node_execution_models.NodeExecution, _artifact_mixin.ExecutionArtifact): - """A class encapsulating a node execution being run on a Flyte remote backend.""" - - def __init__(self, *args, **kwargs): - super(FlyteNodeExecution, self).__init__(*args, **kwargs) - self._task_executions = None - self._subworkflow_node_executions = None - self._inputs = None - self._outputs = None - self._interface = None - - @property - def task_executions(self) -> List["flytekit.remote.tasks.executions.FlyteTaskExecution"]: - return self._task_executions or [] - - @property - def subworkflow_node_executions(self) -> Dict[str, "flytekit.remote.nodes.FlyteNodeExecution"]: - return ( - {} - if self._subworkflow_node_executions is None - else {n.id.node_id: n for n in self._subworkflow_node_executions} - ) - - @property - def executions(self) -> List[_artifact_mixin.ExecutionArtifact]: - return self.task_executions or self._subworkflow_node_executions or [] - - @property - def inputs(self) -> Dict[str, Any]: - """ - Returns the inputs to the execution in the standard python format as dictated by the type engine. - """ - return self._inputs - - @property - def outputs(self) -> Dict[str, Any]: - """ - Returns the outputs to the execution in the standard python format as dictated by the type engine. - - :raises: ``FlyteAssertion`` error if execution is in progress or execution ended in error. - """ - if not self.is_complete: - raise _user_exceptions.FlyteAssertion( - "Please wait until the node execution has completed before requesting the outputs." - ) - if self.error: - raise _user_exceptions.FlyteAssertion("Outputs could not be found because the execution ended in failure.") - return self._outputs - - @property - def error(self) -> _execution_models.ExecutionError: - """ - If execution is in progress, raise an exception. Otherwise, return None if no error was present upon - reaching completion. - """ - if not self.is_complete: - raise _user_exceptions.FlyteAssertion( - "Please wait until the node execution has completed before requesting error information." - ) - return self.closure.error - - @property - def is_complete(self) -> bool: - """Whether or not the execution is complete.""" - return self.closure.phase in { - _execution_models.NodeExecutionPhase.ABORTED, - _execution_models.NodeExecutionPhase.FAILED, - _execution_models.NodeExecutionPhase.SKIPPED, - _execution_models.NodeExecutionPhase.SUCCEEDED, - _execution_models.NodeExecutionPhase.TIMED_OUT, - } - - @classmethod - def promote_from_model(cls, base_model: _node_execution_models.NodeExecution) -> "FlyteNodeExecution": - return cls( - closure=base_model.closure, id=base_model.id, input_uri=base_model.input_uri, metadata=base_model.metadata - ) - - @property - def interface(self) -> "flytekit.remote.interface.TypedInterface": - """ - Return the interface of the task or subworkflow associated with this node execution. - """ - return self._interface - - def sync(self): - """ - Syncs the state of the underlying execution artifact with the state observed by the platform. - """ - if self.metadata.is_parent_node: - if not self.is_complete or self._subworkflow_node_executions is None: - self._subworkflow_node_executions = [ - FlyteNodeExecution.promote_from_model(n) - for n in iterate_node_executions( - _flyte_engine.get_client(), - workflow_execution_identifier=self.id.execution_id, - unique_parent_id=self.id.node_id, - ) - ] - else: - if not self.is_complete or self._task_executions is None: - self._task_executions = [ - FlyteTaskExecution.promote_from_model(t) - for t in iterate_task_executions(_flyte_engine.get_client(), self.id) - ] - - self._sync_closure() - for execution in self.executions: - execution.sync() - - def _sync_closure(self): - """ - Syncs the closure of the underlying execution artifact with the state observed by the platform. - """ - self._closure = _flyte_engine.get_client().get_node_execution(self.id).closure diff --git a/flytekit/remote/remote.py b/flytekit/remote/remote.py index e859db730d..04e18d6794 100644 --- a/flytekit/remote/remote.py +++ b/flytekit/remote/remote.py @@ -1,4 +1,8 @@ -"""Module defining main Flyte backend entrypoint.""" +""" +This module provides the ``FlyteRemote`` object, which is the end-user's main starting point for interacting +with a Flyte backend in an interactive and programmatic way. This of this experience as kind of like the web UI +but in Python object form. +""" from __future__ import annotations import logging @@ -50,7 +54,8 @@ from flytekit.models import launch_plan as launch_plan_models from flytekit.models import literals as literal_models from flytekit.models.admin.common import Sort -from flytekit.models.core.identifier import ResourceType +from flytekit.models.core.identifier import Identifier, ResourceType, WorkflowExecutionIdentifier +from flytekit.models.core.workflow import NodeMetadata from flytekit.models.execution import ( ExecutionMetadata, ExecutionSpec, @@ -58,14 +63,11 @@ NotificationList, WorkflowExecutionGetDataResponse, ) -from flytekit.remote.identifier import Identifier, WorkflowExecutionIdentifier -from flytekit.remote.interface import TypedInterface +from flytekit.remote.executions import FlyteNodeExecution, FlyteTaskExecution, FlyteWorkflowExecution from flytekit.remote.launch_plan import FlyteLaunchPlan -from flytekit.remote.nodes import FlyteNodeExecution -from flytekit.remote.tasks.executions import FlyteTaskExecution -from flytekit.remote.tasks.task import FlyteTask +from flytekit.remote.nodes import FlyteNode +from flytekit.remote.task import FlyteTask from flytekit.remote.workflow import FlyteWorkflow -from flytekit.remote.workflow_execution import FlyteWorkflowExecution ExecutionDataResponse = typing.Union[WorkflowExecutionGetDataResponse, NodeExecutionGetDataResponse] @@ -378,8 +380,6 @@ def fetch_workflow( :param domain: fetch entity from this domain. If None, uses the default_domain attribute. :param name: fetch entity with matching name. :param version: fetch entity with matching version. If None, gets the latest version of the entity. - :returns: :class:`~flytekit.remote.workflow.FlyteWorkflow` - :raises: FlyteAssertion if name is None """ if name is None: @@ -395,26 +395,15 @@ def fetch_workflow( admin_workflow = self.client.get_workflow(workflow_id) compiled_wf = admin_workflow.closure.compiled_workflow - base_model = compiled_wf.primary.template - sub_workflows = {sw.template.id: sw.template for sw in compiled_wf.sub_workflows} - tasks = {t.template.id: t.template for t in compiled_wf.tasks} - node_launch_plans = {} # TODO: Inspect branch nodes for launch plans - for node in FlyteWorkflow.get_non_system_nodes(base_model.nodes): + for node in FlyteWorkflow.get_non_system_nodes(compiled_wf.primary.template.nodes): if node.workflow_node is not None and node.workflow_node.launchplan_ref is not None: node_launch_plans[node.workflow_node.launchplan_ref] = self.client.get_launch_plan( node.workflow_node.launchplan_ref ).spec - flyte_workflow = FlyteWorkflow.promote_from_model( - base_model=compiled_wf.primary.template, - sub_workflows=sub_workflows, - node_launch_plans=node_launch_plans, - tasks=tasks, - ) - flyte_workflow._id = workflow_id - return flyte_workflow + return FlyteWorkflow.promote_from_closure(compiled_wf, node_launch_plans) def fetch_launch_plan( self, project: str = None, domain: str = None, name: str = None, version: str = None @@ -1015,7 +1004,7 @@ def sync( self, execution: FlyteWorkflowExecution, entity_definition: typing.Union[FlyteWorkflow, FlyteTask] = None, - sync_nodes: bool = True, + sync_nodes: bool = False, ) -> FlyteWorkflowExecution: """ This function was previously a singledispatchmethod. We've removed that but this function remains @@ -1036,63 +1025,198 @@ def sync_workflow_execution( self, execution: FlyteWorkflowExecution, entity_definition: typing.Union[FlyteWorkflow, FlyteTask] = None, - sync_nodes: bool = True, + sync_nodes: bool = False, ) -> FlyteWorkflowExecution: - - """Sync a FlyteWorkflowExecution object with its corresponding remote state.""" + """ + Sync a FlyteWorkflowExecution object with its corresponding remote state. + """ if entity_definition is not None: raise ValueError("Entity definition arguments aren't supported when syncing workflow executions") + + # Update closure, and then data, because we don't want the execution to finish between when we get the data, + # and then for the closure to have is_complete to be true. + execution._closure = self.client.get_execution(execution.id).closure execution_data = self.client.get_execution_data(execution.id) lp_id = execution.spec.launch_plan + if sync_nodes: + underlying_node_executions = [ + FlyteNodeExecution.promote_from_model(n) for n in iterate_node_executions(self.client, execution.id) + ] + if execution.spec.launch_plan.resource_type == ResourceType.TASK: + # This condition is only true for single-task executions flyte_entity = self.fetch_task(lp_id.project, lp_id.domain, lp_id.name, lp_id.version) + if sync_nodes: + # Need to construct the mapping. There should've been returned exactly three nodes, a start, + # an end, and a task node. + task_node_exec = [ + x + for x in filter( + lambda x: x.id.node_id != constants.START_NODE_ID and x.id.node_id != constants.END_NODE_ID, + underlying_node_executions, + ) + ] + # We need to manually make a map of the nodes since there is none for single task executions + # Assume the first one is the only one. + node_mapping = ( + { + task_node_exec[0].id.node_id: FlyteNode( + id=flyte_entity.id, + upstream_nodes=[], + bindings=[], + metadata=NodeMetadata(name=""), + flyte_task=flyte_entity, + ) + } + if len(task_node_exec) >= 1 + else {} # This is for the case where node executions haven't appeared yet + ) else: + # This is the default case, an execution of a normal workflow through a launch plan wf_id = self.fetch_launch_plan(lp_id.project, lp_id.domain, lp_id.name, lp_id.version).workflow_id flyte_entity = self.fetch_workflow(wf_id.project, wf_id.domain, wf_id.name, wf_id.version) + execution._flyte_workflow = flyte_entity + node_mapping = flyte_entity._node_map - # sync closure, node executions, and inputs/outputs - execution._closure = self.client.get_execution(execution.id).closure + # update node executions (if requested), and inputs/outputs if sync_nodes: - execution._node_executions = { - node.id.node_id: self.sync_node_execution(FlyteNodeExecution.promote_from_model(node), flyte_entity) - for node in iterate_node_executions(self.client, execution.id) - } + node_execs = {} + for n in underlying_node_executions: + node_execs[n.id.node_id] = self.sync_node_execution(n, node_mapping) + execution._node_executions = node_execs return self._assign_inputs_and_outputs(execution, execution_data, flyte_entity.interface) def sync_node_execution( - self, execution: FlyteNodeExecution, entity_definition: typing.Union[FlyteWorkflow, FlyteTask] = None + self, execution: FlyteNodeExecution, node_mapping: typing.Dict[str, FlyteNode] ) -> FlyteNodeExecution: - """Sync a FlyteNodeExecution object with its corresponding remote state.""" - if ( - execution.id.node_id in {constants.START_NODE_ID, constants.END_NODE_ID} - or execution.id.node_id.endswith(constants.START_NODE_ID) - or execution.id.node_id.endswith(constants.END_NODE_ID) - ): + """ + Get data backing a node execution. These FlyteNodeExecution objects should've come from Admin with the model + fields already populated correctly. For purposes of the remote experience, we'd like to supplement the object + with some additional fields: + - inputs/outputs + - task/workflow executions, and/or underlying node executions in the case of parent nodes + - TypedInterface (remote wrapper type) + + A node can have several different types of executions behind it. That is, the node could've run (perhaps + multiple times because of retries): + - A task + - A static subworkflow + - A dynamic subworkflow (which in turn may have run additional tasks, subwfs, and/or launch plans) + - A launch plan + + The data model is complicated, so ascertaining which of these happened is a bit tricky. That logic is + encapsulated in this function. + """ + # For single task execution - the metadata spec node id is missing. In these cases, revert to regular node id + node_id = execution.metadata.spec_node_id + if not node_id: + node_id = execution.id.node_id + remote_logger.debug(f"No metadata spec_node_id found, using {node_id}") + + # First see if it's a dummy node, if it is, we just skip it. + if constants.START_NODE_ID in node_id or constants.END_NODE_ID in node_id: return execution - # sync closure, child nodes, interface, and inputs/outputs - execution._closure = self.client.get_node_execution(execution.id).closure + # Look for the Node object in the mapping supplied + if node_id in node_mapping: + execution._node = node_mapping[node_id] + else: + raise Exception(f"Missing node from mapping: {node_id}") + + # Get the node execution data + node_execution_get_data_response = self.client.get_node_execution_data(execution.id) + + # Calling a launch plan directly case + # If a node ran a launch plan directly (i.e. not through a dynamic task or anything) then + # the closure should have a workflow_node_metadata populated with the launched execution id. + # The parent node flag should not be populated here + # This is the simplest case + if not execution.metadata.is_parent_node and execution.closure.workflow_node_metadata: + launched_exec_id = execution.closure.workflow_node_metadata.execution_id + # This is a recursive call, basically going through the same process that brought us here in the first + # place, but on the launched execution. + launched_exec = self.fetch_workflow_execution( + project=launched_exec_id.project, domain=launched_exec_id.domain, name=launched_exec_id.name + ) + self.sync_workflow_execution(launched_exec) + if launched_exec.is_complete: + # The synced underlying execution should've had these populated. + execution._inputs = launched_exec.inputs + execution._outputs = launched_exec.outputs + execution._workflow_executions.append(launched_exec) + execution._interface = launched_exec._flyte_workflow.interface + return execution + + # If a node ran a static subworkflow or a dynamic subworkflow then the parent flag will be set. if execution.metadata.is_parent_node: - execution._subworkflow_node_executions = [ - self.sync_node_execution(FlyteNodeExecution.promote_from_model(node), entity_definition) - for node in iterate_node_executions( - self.client, - workflow_execution_identifier=execution.id.execution_id, - unique_parent_id=execution.id.node_id, - ) - ] + # We'll need to query child node executions regardless since this is a parent node + child_node_executions = iterate_node_executions( + self.client, + workflow_execution_identifier=execution.id.execution_id, + unique_parent_id=execution.id.node_id, + ) + child_node_executions = [x for x in child_node_executions] + + # If this was a dynamic task, then there should be a CompiledWorkflowClosure inside the + # NodeExecutionGetDataResponse + if node_execution_get_data_response.dynamic_workflow is not None: + compiled_wf = node_execution_get_data_response.dynamic_workflow.compiled_workflow + node_launch_plans = {} + # TODO: Inspect branch nodes for launch plans + for node in FlyteWorkflow.get_non_system_nodes(compiled_wf.primary.template.nodes): + if ( + node.workflow_node is not None + and node.workflow_node.launchplan_ref is not None + and node.workflow_node.launchplan_ref not in node_launch_plans + ): + node_launch_plans[node.workflow_node.launchplan_ref] = self.client.get_launch_plan( + node.workflow_node.launchplan_ref + ).spec + + dynamic_flyte_wf = FlyteWorkflow.promote_from_closure(compiled_wf, node_launch_plans) + execution._underlying_node_executions = [ + self.sync_node_execution(FlyteNodeExecution.promote_from_model(cne), dynamic_flyte_wf._node_map) + for cne in child_node_executions + ] + # This is copied from below - dynamic tasks have both task executions (executions of the parent + # task) as well as underlying node executions (of the generated subworkflow). Feel free to refactor + # if you can think of a better way. + execution._task_executions = [ + self.sync_task_execution(FlyteTaskExecution.promote_from_model(t)) + for t in iterate_task_executions(self.client, execution.id) + ] + execution._interface = dynamic_flyte_wf.interface + else: + # If it does not, then it should be a static subworkflow + if not isinstance(execution._node.flyte_entity, FlyteWorkflow): + remote_logger.error( + f"NE {execution} entity should be a workflow, {type(execution._node)}, {execution._node}" + ) + raise Exception(f"Node entity has type {type(execution._node)}") + sub_flyte_workflow = execution._node.flyte_entity + sub_node_mapping = {n.id: n for n in sub_flyte_workflow.flyte_nodes} + execution._underlying_node_executions = [ + self.sync_node_execution(FlyteNodeExecution.promote_from_model(cne), sub_node_mapping) + for cne in child_node_executions + ] + execution._interface = sub_flyte_workflow.interface + + # This is the plain ol' task execution case else: execution._task_executions = [ self.sync_task_execution(FlyteTaskExecution.promote_from_model(t)) for t in iterate_task_executions(self.client, execution.id) ] - execution._interface = self._get_node_execution_interface(execution, entity_definition) - return self._assign_inputs_and_outputs( + execution._interface = execution._node.flyte_entity.interface + + self._assign_inputs_and_outputs( execution, - self.client.get_node_execution_data(execution.id), + node_execution_get_data_response, execution.interface, ) + return execution + def sync_task_execution( self, execution: FlyteTaskExecution, entity_definition: typing.Union[FlyteWorkflow, FlyteTask] = None ) -> FlyteTaskExecution: @@ -1123,7 +1247,12 @@ def terminate(self, execution: FlyteWorkflowExecution, cause: str): # Helper Methods # ################## - def _assign_inputs_and_outputs(self, execution, execution_data, interface): + def _assign_inputs_and_outputs( + self, + execution: typing.Union[FlyteWorkflowExecution, FlyteNodeExecution, FlyteTaskExecution], + execution_data, + interface, + ): """Helper for assigning synced inputs and outputs to an execution object.""" with self.remote_context() as ctx: execution._inputs = TypeEngine.literal_map_to_kwargs( @@ -1164,49 +1293,3 @@ def _get_output_literal_map(self, execution_data: ExecutionDataResponse) -> lite common_utils.load_proto_from_file(literals_pb2.LiteralMap, tmp_name) ) return literal_models.LiteralMap({}) - - def _get_node_execution_interface( - self, node_execution: FlyteNodeExecution, entity_definition: typing.Union[FlyteWorkflow, FlyteTask] - ) -> TypedInterface: - """Return the interface of the task or subworkflow associated with this node execution.""" - if isinstance(entity_definition, FlyteTask): - # A single task execution consists of a Flyte workflow with single node whose interface matches that of - # the underlying task - return entity_definition.interface - - for node in entity_definition.flyte_nodes: - if node.id == node_execution.id.node_id: - if node.task_node is not None: - return node.task_node.flyte_task.interface - elif node.workflow_node is not None and node.workflow_node.sub_workflow_ref is not None: - # Fetch the workflow and use its interface - sub_workflow_ref = node.workflow_node.sub_workflow_ref - workflow = self.fetch_workflow( - sub_workflow_ref.project, - sub_workflow_ref.domain, - sub_workflow_ref.name, - sub_workflow_ref.version, - ) - return workflow.interface - elif node.workflow_node is not None and node.workflow_node.launchplan_ref is not None: - # Fetch the launch plan this node launched, and from there fetch the referenced workflow and use its - # interface. - lp_ref = node.workflow_node.launchplan_ref - launch_plan = self.fetch_launch_plan(lp_ref.project, lp_ref.domain, lp_ref.name, lp_ref.version) - workflow = self.fetch_workflow( - launch_plan.workflow_id.project, - launch_plan.workflow_id.domain, - launch_plan.workflow_id.name, - launch_plan.workflow_id.version, - ) - return workflow.interface - - # dynamically generated nodes won't have a corresponding node in the compiled workflow closure. - # in that case, we fetch the interface from the underlying task execution they ran - if len(node_execution.task_executions) > 0: - # if not a parent node, assume a task execution node - task_id = node_execution.task_executions[0].id.task_id - task = self.fetch_task(task_id.project, task_id.domain, task_id.name, task_id.version) - return task.interface - - remote_logger.info("failed to find node interface from entity definition closure") diff --git a/flytekit/remote/tasks/task.py b/flytekit/remote/task.py similarity index 95% rename from flytekit/remote/tasks/task.py rename to flytekit/remote/task.py index 967f4e43aa..0c48f5f15e 100644 --- a/flytekit/remote/tasks/task.py +++ b/flytekit/remote/task.py @@ -6,7 +6,6 @@ from flytekit.loggers import logger from flytekit.models import task as _task_model from flytekit.models.core import identifier as _identifier_model -from flytekit.remote import identifier as _identifier from flytekit.remote import interface as _interfaces @@ -61,7 +60,7 @@ def promote_from_model(cls, base_model: _task_model.TaskTemplate) -> "FlyteTask" ) # Override the newly generated name if one exists in the base model if not base_model.id.is_empty: - t._id = _identifier.Identifier.promote_from_model(base_model.id) + t._id = base_model.id if t.interface is not None: try: diff --git a/flytekit/remote/tasks/__init__.py b/flytekit/remote/tasks/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/flytekit/remote/tasks/executions.py b/flytekit/remote/tasks/executions.py deleted file mode 100644 index 9937b4be77..0000000000 --- a/flytekit/remote/tasks/executions.py +++ /dev/null @@ -1,96 +0,0 @@ -from typing import Any, Dict, Optional - -from flytekit.clients.helpers import iterate_node_executions as _iterate_node_executions -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.mixins import artifact as _artifact_mixin -from flytekit.engines.flyte import engine as _flyte_engine -from flytekit.models.admin import task_execution as _task_execution_model -from flytekit.models.core import execution as _execution_models - - -class FlyteTaskExecution(_task_execution_model.TaskExecution, _artifact_mixin.ExecutionArtifact): - """A class encapsulating a task execution being run on a Flyte remote backend.""" - - def __init__(self, *args, **kwargs): - super(FlyteTaskExecution, self).__init__(*args, **kwargs) - self._inputs = None - self._outputs = None - - @property - def is_complete(self) -> bool: - """Whether or not the execution is complete.""" - return self.closure.phase in { - _execution_models.TaskExecutionPhase.ABORTED, - _execution_models.TaskExecutionPhase.FAILED, - _execution_models.TaskExecutionPhase.SUCCEEDED, - } - - @property - def inputs(self) -> Dict[str, Any]: - """ - Returns the inputs of the task execution in the standard Python format that is produced by - the type engine. - """ - return self._inputs - - @property - def outputs(self) -> Dict[str, Any]: - """ - Returns the outputs of the task execution, if available, in the standard Python format that is produced by - the type engine. - - :raises: ``FlyteAssertion`` error if execution is in progress or execution ended in error. - """ - if not self.is_complete: - raise _user_exceptions.FlyteAssertion( - "Please wait until the node execution has completed before requesting the outputs." - ) - if self.error: - raise _user_exceptions.FlyteAssertion("Outputs could not be found because the execution ended in failure.") - return self._outputs - - @property - def error(self) -> Optional[_execution_models.ExecutionError]: - """ - If execution is in progress, raise an exception. Otherwise, return None if no error was present upon - reaching completion. - """ - if not self.is_complete: - raise _user_exceptions.FlyteAssertion( - "Please what until the task execution has completed before requesting error information." - ) - return self.closure.error - - def get_child_executions(self, filters=None): - from flytekit.remote import nodes as _nodes - - if not self.is_parent: - raise _user_exceptions.FlyteAssertion("Only task executions marked with 'is_parent' have child executions.") - client = _flyte_engine.get_client() - models = { - v.id.node_id: v - for v in _iterate_node_executions(client, task_execution_identifier=self.id, filters=filters) - } - - return {k: _nodes.FlyteNodeExecution.promote_from_model(v) for k, v in models.items()} - - @classmethod - def promote_from_model(cls, base_model: _task_execution_model.TaskExecution) -> "FlyteTaskExecution": - return cls( - closure=base_model.closure, - id=base_model.id, - input_uri=base_model.input_uri, - is_parent=base_model.is_parent, - ) - - def sync(self): - """ - Syncs the state of the underlying execution artifact with the state observed by the platform. - """ - self._sync_closure() - - def _sync_closure(self): - """ - Syncs the closure of the underlying execution artifact with the state observed by the platform. - """ - self._closure = _flyte_engine.get_client().get_task_execution(self.id).closure diff --git a/flytekit/remote/workflow.py b/flytekit/remote/workflow.py index 66bb7f00e7..396f377500 100644 --- a/flytekit/remote/workflow.py +++ b/flytekit/remote/workflow.py @@ -1,16 +1,17 @@ +from __future__ import annotations + from typing import Dict, List, Optional from flytekit.common import constants as _constants -from flytekit.common.exceptions import system as _system_exceptions from flytekit.common.exceptions import user as _user_exceptions from flytekit.common.mixins import hash as _hash_mixin from flytekit.core.interface import Interface from flytekit.core.type_engine import TypeEngine -from flytekit.models import launch_plan as _launch_plan_models +from flytekit.models import launch_plan as launch_plan_models from flytekit.models import task as _task_models -from flytekit.models.core import identifier as _identifier_model +from flytekit.models.core import compiler as compiler_models +from flytekit.models.core import identifier as id_models from flytekit.models.core import workflow as _workflow_models -from flytekit.remote import identifier as _identifier from flytekit.remote import interface as _interfaces from flytekit.remote import nodes as _nodes @@ -23,10 +24,15 @@ def __init__( nodes: List[_nodes.FlyteNode], interface, output_bindings, - id, + id: id_models.Identifier, metadata, metadata_defaults, + subworkflows: Optional[Dict[id_models.Identifier, _workflow_models.WorkflowTemplate]] = None, + tasks: Optional[Dict[id_models.Identifier, _task_models.TaskSpec]] = None, + launch_plans: Optional[Dict[id_models.Identifier, launch_plan_models.LaunchPlanSpec]] = None, + compiled_closure: Optional[compiler_models.CompiledWorkflowClosure] = None, ): + # TODO: Remove check for node in nodes: for upstream in node.upstream_nodes: if upstream.id is None: @@ -46,9 +52,13 @@ def __init__( self._flyte_nodes = nodes self._python_interface = None - @property - def upstream_entities(self): - return set(n.executable_flyte_object for n in self._flyte_nodes) + # Optional things that we save for ease of access when promoting from a model or CompiledWorkflowClosure + self._subworkflows = subworkflows + self._tasks = tasks + self._launch_plans = launch_plans + self._compiled_closure = compiled_closure + + self._node_map = None @property def interface(self) -> _interfaces.TypedInterface: @@ -60,7 +70,7 @@ def entity_type_text(self) -> str: @property def resource_type(self): - return _identifier_model.ResourceType.WORKFLOW + return id_models.ResourceType.WORKFLOW @property def flyte_nodes(self) -> List[_nodes.FlyteNode]: @@ -76,37 +86,6 @@ def guessed_python_interface(self, value): return self._python_interface = value - def get_sub_workflows(self) -> List["FlyteWorkflow"]: - result = [] - for node in self.flyte_nodes: - if node.workflow_node is not None and node.workflow_node.sub_workflow_ref is not None: - if node.flyte_entity is not None and node.flyte_entity.entity_type_text == "Workflow": - result.append(node.flyte_entity) - result.extend(node.flyte_entity.get_sub_workflows()) - else: - raise _system_exceptions.FlyteSystemException( - "workflow node with subworkflow found but bad executable " "object {}".format(node.flyte_entity) - ) - - # get subworkflows in conditional branches - if node.branch_node is not None: - if_else: _workflow_models.IfElseBlock = node.branch_node.if_else - leaf_nodes: List[_nodes.FlyteNode] = filter( - None, - [ - if_else.case.then_node, - *([] if if_else.other is None else [x.then_node for x in if_else.other]), - if_else.else_node, - ], - ) - for leaf_node in leaf_nodes: - exec_flyte_obj = leaf_node.flyte_entity - if exec_flyte_obj is not None and exec_flyte_obj.entity_type_text == "Workflow": - result.append(exec_flyte_obj) - result.extend(exec_flyte_obj.get_sub_workflows()) - - return result - @classmethod def get_non_system_nodes(cls, nodes: List[_workflow_models.Node]) -> List[_workflow_models.Node]: return [n for n in nodes if n.id not in {_constants.START_NODE_ID, _constants.END_NODE_ID}] @@ -115,10 +94,10 @@ def get_non_system_nodes(cls, nodes: List[_workflow_models.Node]) -> List[_workf def promote_from_model( cls, base_model: _workflow_models.WorkflowTemplate, - sub_workflows: Optional[Dict[_identifier.Identifier, _workflow_models.WorkflowTemplate]] = None, - node_launch_plans: Optional[Dict[_identifier.Identifier, _launch_plan_models.LaunchPlanSpec]] = None, - tasks: Optional[Dict[_identifier.Identifier, _task_models.TaskTemplate]] = None, - ) -> "FlyteWorkflow": + sub_workflows: Optional[Dict[id_models, _workflow_models.WorkflowTemplate]] = None, + node_launch_plans: Optional[Dict[id_models, launch_plan_models.LaunchPlanSpec]] = None, + tasks: Optional[Dict[id_models, _task_models.TaskTemplate]] = None, + ) -> FlyteWorkflow: base_model_non_system_nodes = cls.get_non_system_nodes(base_model.nodes) sub_workflows = sub_workflows or {} tasks = tasks or {} @@ -137,11 +116,14 @@ def promote_from_model( # No inputs/outputs specified, see the constructor for more information on the overrides. wf = cls( nodes=list(node_map.values()), - id=_identifier.Identifier.promote_from_model(base_model.id), + id=base_model.id, metadata=base_model.metadata, metadata_defaults=base_model.metadata_defaults, interface=_interfaces.TypedInterface.promote_from_model(base_model.interface), output_bindings=base_model.outputs, + subworkflows=sub_workflows, + tasks=tasks, + launch_plans=node_launch_plans, ) if wf.interface is not None: @@ -149,8 +131,35 @@ def promote_from_model( inputs=TypeEngine.guess_python_types(wf.interface.inputs), outputs=TypeEngine.guess_python_types(wf.interface.outputs), ) + wf._node_map = node_map return wf + @classmethod + def promote_from_closure( + cls, + closure: compiler_models.CompiledWorkflowClosure, + node_launch_plans: Optional[Dict[id_models, launch_plan_models.LaunchPlanSpec]] = None, + ): + """ + Extracts out the relevant portions of a FlyteWorkflow from a closure from the control plane. + + :param closure: This is the closure returned by Admin + :param node_launch_plans: The reason this exists is because the compiled closure doesn't have launch plans. + It only has subworkflows and tasks. Why this is is unclear. If supplied, this map of launch plans will be + :return: + """ + sub_workflows = {sw.template.id: sw.template for sw in closure.sub_workflows} + tasks = {t.template.id: t.template for t in closure.tasks} + + flyte_wf = FlyteWorkflow.promote_from_model( + base_model=closure.primary.template, + sub_workflows=sub_workflows, + node_launch_plans=node_launch_plans, + tasks=tasks, + ) + flyte_wf._compiled_closure = closure + return flyte_wf + def __call__(self, *args, **input_map): raise NotImplementedError diff --git a/flytekit/remote/workflow_execution.py b/flytekit/remote/workflow_execution.py index 0c201c9056..e69de29bb2 100644 --- a/flytekit/remote/workflow_execution.py +++ b/flytekit/remote/workflow_execution.py @@ -1,76 +0,0 @@ -from typing import Any, Dict - -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.models import execution as _execution_models -from flytekit.models.core import execution as _core_execution_models -from flytekit.remote import identifier as _core_identifier -from flytekit.remote import nodes as _nodes - - -class FlyteWorkflowExecution(_execution_models.Execution): - """A class encapsulating a workflow execution being run on a Flyte remote backend.""" - - def __init__(self, *args, **kwargs): - super(FlyteWorkflowExecution, self).__init__(*args, **kwargs) - self._node_executions = None - self._inputs = None - self._outputs = None - - @property - def node_executions(self) -> Dict[str, _nodes.FlyteNodeExecution]: - """Get a dictionary of node executions that are a part of this workflow execution.""" - return self._node_executions or {} - - @property - def inputs(self) -> Dict[str, Any]: - """ - Returns the inputs to the execution in the standard python format as dictated by the type engine. - """ - return self._inputs - - @property - def outputs(self) -> Dict[str, Any]: - """ - Returns the outputs to the execution in the standard python format as dictated by the type engine. - - :raises: ``FlyteAssertion`` error if execution is in progress or execution ended in error. - """ - if not self.is_complete: - raise _user_exceptions.FlyteAssertion( - "Please wait until the node execution has completed before requesting the outputs." - ) - if self.error: - raise _user_exceptions.FlyteAssertion("Outputs could not be found because the execution ended in failure.") - return self._outputs - - @property - def error(self) -> _core_execution_models.ExecutionError: - """ - If execution is in progress, raise an exception. Otherwise, return None if no error was present upon - reaching completion. - """ - if not self.is_complete: - raise _user_exceptions.FlyteAssertion( - "Please wait until a workflow has completed before checking for an error." - ) - return self.closure.error - - @property - def is_complete(self) -> bool: - """ - Whether or not the execution is complete. - """ - return self.closure.phase in { - _core_execution_models.WorkflowExecutionPhase.ABORTED, - _core_execution_models.WorkflowExecutionPhase.FAILED, - _core_execution_models.WorkflowExecutionPhase.SUCCEEDED, - _core_execution_models.WorkflowExecutionPhase.TIMED_OUT, - } - - @classmethod - def promote_from_model(cls, base_model: _execution_models.Execution) -> "FlyteWorkflowExecution": - return cls( - closure=base_model.closure, - id=_core_identifier.WorkflowExecutionIdentifier.promote_from_model(base_model.id), - spec=base_model.spec, - ) diff --git a/tests/flytekit/integration/remote/test_remote.py b/tests/flytekit/integration/remote/test_remote.py index c00564a5f3..9dc32a2e54 100644 --- a/tests/flytekit/integration/remote/test_remote.py +++ b/tests/flytekit/integration/remote/test_remote.py @@ -82,7 +82,7 @@ def test_monitor_workflow_execution(flyteclient, flyte_workflows_register, flyte poll_interval = datetime.timedelta(seconds=1) time_to_give_up = datetime.datetime.utcnow() + datetime.timedelta(seconds=60) - execution = remote.sync_workflow_execution(execution) + execution = remote.sync_workflow_execution(execution, sync_nodes=True) while datetime.datetime.utcnow() < time_to_give_up: if execution.is_complete: @@ -94,7 +94,7 @@ def test_monitor_workflow_execution(flyteclient, flyte_workflows_register, flyte execution.outputs time.sleep(poll_interval.total_seconds()) - execution = remote.sync_workflow_execution(execution) + execution = remote.sync_workflow_execution(execution, sync_nodes=True) if execution.node_executions: assert execution.node_executions["start-node"].closure.phase == 3 # SUCCEEEDED diff --git a/tests/flytekit/unit/models/admin/test_node_executions.py b/tests/flytekit/unit/models/admin/test_node_executions.py index 84d8785d09..b4cd77e5e8 100644 --- a/tests/flytekit/unit/models/admin/test_node_executions.py +++ b/tests/flytekit/unit/models/admin/test_node_executions.py @@ -1,7 +1,51 @@ from flytekit.models import node_execution as node_execution_models +from flytekit.models.core import catalog, identifier +from tests.flytekit.unit.common_tests.test_workflow_promote import get_compiled_workflow_closure def test_metadata(): md = node_execution_models.NodeExecutionMetaData(retry_group="0", is_parent_node=True, spec_node_id="n0") md2 = node_execution_models.NodeExecutionMetaData.from_flyte_idl(md.to_flyte_idl()) assert md == md2 + + +def test_workflow_node_metadata(): + wf_exec_id = identifier.WorkflowExecutionIdentifier("project", "domain", "name") + + obj = node_execution_models.WorkflowNodeMetadata(execution_id=wf_exec_id) + assert obj.execution_id is wf_exec_id + + obj2 = node_execution_models.WorkflowNodeMetadata.from_flyte_idl(obj.to_flyte_idl()) + assert obj == obj2 + + +def test_task_node_metadata(): + task_id = identifier.Identifier(identifier.ResourceType.TASK, "project", "domain", "name", "version") + wf_exec_id = identifier.WorkflowExecutionIdentifier("project", "domain", "name") + node_exec_id = identifier.NodeExecutionIdentifier( + "node_id", + wf_exec_id, + ) + te_id = identifier.TaskExecutionIdentifier(task_id, node_exec_id, 3) + ds_id = identifier.Identifier(identifier.ResourceType.TASK, "project", "domain", "t1", "abcdef") + tag = catalog.CatalogArtifactTag("my-artifact-id", "some name") + catalog_metadata = catalog.CatalogMetadata(dataset_id=ds_id, artifact_tag=tag, source_task_execution=te_id) + + obj = node_execution_models.TaskNodeMetadata(cache_status=0, catalog_key=catalog_metadata) + assert obj.cache_status == 0 + assert obj.catalog_key == catalog_metadata + + obj2 = node_execution_models.TaskNodeMetadata.from_flyte_idl(obj.to_flyte_idl()) + assert obj2 == obj + + +def test_dynamic_wf_node_metadata(): + wf_id = identifier.Identifier(identifier.ResourceType.WORKFLOW, "project", "domain", "name", "version") + cwc = get_compiled_workflow_closure() + + obj = node_execution_models.DynamicWorkflowNodeMetadata(id=wf_id, compiled_workflow=cwc) + assert obj.id == wf_id + assert obj.compiled_workflow == cwc + + obj2 = node_execution_models.DynamicWorkflowNodeMetadata.from_flyte_idl(obj.to_flyte_idl()) + assert obj2 == obj diff --git a/tests/flytekit/unit/models/core/test_catalog.py b/tests/flytekit/unit/models/core/test_catalog.py new file mode 100644 index 0000000000..c89fe6545c --- /dev/null +++ b/tests/flytekit/unit/models/core/test_catalog.py @@ -0,0 +1,32 @@ +from flytekit.models.core import catalog, identifier + + +def test_catalog_artifact_tag(): + obj = catalog.CatalogArtifactTag("my-artifact-id", "some name") + assert obj.artifact_id == "my-artifact-id" + assert obj.name == "some name" + + obj2 = catalog.CatalogArtifactTag.from_flyte_idl(obj.to_flyte_idl()) + assert obj == obj2 + assert obj2.artifact_id == "my-artifact-id" + assert obj2.name == "some name" + + +def test_catalog_metadata(): + task_id = identifier.Identifier(identifier.ResourceType.TASK, "project", "domain", "name", "version") + wf_exec_id = identifier.WorkflowExecutionIdentifier("project", "domain", "name") + node_exec_id = identifier.NodeExecutionIdentifier( + "node_id", + wf_exec_id, + ) + te_id = identifier.TaskExecutionIdentifier(task_id, node_exec_id, 3) + ds_id = identifier.Identifier(identifier.ResourceType.TASK, "project", "domain", "t1", "abcdef") + tag = catalog.CatalogArtifactTag("my-artifact-id", "some name") + obj = catalog.CatalogMetadata(dataset_id=ds_id, artifact_tag=tag, source_task_execution=te_id) + assert obj.dataset_id is ds_id + assert obj.source_execution is te_id + assert obj.source_task_execution is te_id + assert obj.artifact_tag is tag + + obj2 = catalog.CatalogMetadata.from_flyte_idl(obj.to_flyte_idl()) + assert obj == obj2 diff --git a/tests/flytekit/unit/remote/test_identifier.py b/tests/flytekit/unit/remote/test_identifier.py deleted file mode 100644 index 0335b1db76..0000000000 --- a/tests/flytekit/unit/remote/test_identifier.py +++ /dev/null @@ -1,77 +0,0 @@ -import pytest - -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.models.core import identifier as _core_identifier -from flytekit.remote import identifier as _identifier - - -def test_identifier(): - identifier = _identifier.Identifier(_core_identifier.ResourceType.WORKFLOW, "project", "domain", "name", "v1") - assert identifier == _identifier.Identifier.from_urn("wf:project:domain:name:v1") - assert identifier == _core_identifier.Identifier( - _core_identifier.ResourceType.WORKFLOW, "project", "domain", "name", "v1" - ) - assert identifier.__str__() == "wf:project:domain:name:v1" - - -@pytest.mark.parametrize( - "urn", - [ - "", - "project:domain:name:v1", - "wf:project:domain:name:v1:foobar", - "foobar:project:domain:name:v1", - ], -) -def test_identifier_exceptions(urn): - with pytest.raises(_user_exceptions.FlyteValueException): - _identifier.Identifier.from_urn(urn) - - -def test_workflow_execution_identifier(): - identifier = _identifier.WorkflowExecutionIdentifier("project", "domain", "name") - assert identifier == _identifier.WorkflowExecutionIdentifier.from_urn("ex:project:domain:name") - assert identifier == _identifier.WorkflowExecutionIdentifier.promote_from_model( - _core_identifier.WorkflowExecutionIdentifier("project", "domain", "name") - ) - assert identifier.__str__() == "ex:project:domain:name" - - -@pytest.mark.parametrize( - "urn", ["", "project:domain:name", "project:domain:name:foobar", "ex:project:domain:name:foobar"] -) -def test_workflow_execution_identifier_exceptions(urn): - with pytest.raises(_user_exceptions.FlyteValueException): - _identifier.WorkflowExecutionIdentifier.from_urn(urn) - - -def test_task_execution_identifier(): - task_id = _identifier.Identifier(_core_identifier.ResourceType.TASK, "project", "domain", "name", "version") - node_execution_id = _core_identifier.NodeExecutionIdentifier( - node_id="n0", execution_id=_core_identifier.WorkflowExecutionIdentifier("project", "domain", "name") - ) - identifier = _identifier.TaskExecutionIdentifier( - task_id=task_id, - node_execution_id=node_execution_id, - retry_attempt=0, - ) - assert identifier == _identifier.TaskExecutionIdentifier.from_urn( - "te:project:domain:name:n0:project:domain:name:version:0" - ) - assert identifier == _identifier.TaskExecutionIdentifier.promote_from_model( - _core_identifier.TaskExecutionIdentifier(task_id, node_execution_id, 0) - ) - assert identifier.__str__() == "te:project:domain:name:n0:project:domain:name:version:0" - - -@pytest.mark.parametrize( - "urn", - [ - "", - "te:project:domain:name:n0:project:domain:name:version", - "foobar:project:domain:name:n0:project:domain:name:version:0", - ], -) -def test_task_execution_identifier_exceptions(urn): - with pytest.raises(_user_exceptions.FlyteValueException): - _identifier.TaskExecutionIdentifier.from_urn(urn) diff --git a/tests/flytekit/unit/remote/test_remote.py b/tests/flytekit/unit/remote/test_remote.py index 873f67d752..2bb8174069 100644 --- a/tests/flytekit/unit/remote/test_remote.py +++ b/tests/flytekit/unit/remote/test_remote.py @@ -6,20 +6,8 @@ from flytekit.common.exceptions import user as user_exceptions from flytekit.configuration import internal from flytekit.models import common as common_models -from flytekit.models.admin.workflow import Workflow -from flytekit.models.core.identifier import ( - Identifier, - NodeExecutionIdentifier, - ResourceType, - WorkflowExecutionIdentifier, -) +from flytekit.models.core.identifier import ResourceType, WorkflowExecutionIdentifier from flytekit.models.execution import Execution -from flytekit.models.interface import TypedInterface, Variable -from flytekit.models.launch_plan import LaunchPlan -from flytekit.models.node_execution import NodeExecution, NodeExecutionMetaData -from flytekit.models.task import Task -from flytekit.models.types import LiteralType, SimpleType -from flytekit.remote import FlyteWorkflow from flytekit.remote.remote import FlyteRemote CLIENT_METHODS = { @@ -41,50 +29,6 @@ } -@patch("flytekit.clients.friendly.SynchronousFlyteClient") -@patch("flytekit.configuration.platform.URL") -@patch("flytekit.configuration.platform.INSECURE") -@pytest.mark.parametrize( - "entity_cls,resource_type", - [ - [Workflow, ResourceType.WORKFLOW], - [Task, ResourceType.TASK], - [LaunchPlan, ResourceType.LAUNCH_PLAN], - ], -) -def test_remote_fetch_execute_entities_task_workflow_launchplan( - mock_insecure, - mock_url, - mock_client, - entity_cls, - resource_type, -): - admin_entities = [ - entity_cls( - Identifier(resource_type, "p1", "d1", "n1", version), - *([MagicMock()] if resource_type != ResourceType.LAUNCH_PLAN else [MagicMock(), MagicMock()]), - ) - for version in ["latest", "old"] - ] - - mock_url.get.return_value = "localhost" - mock_insecure.get.return_value = True - mock_client = MagicMock() - getattr(mock_client, CLIENT_METHODS[resource_type]).return_value = admin_entities, "" - - remote = FlyteRemote.from_config("p1", "d1") - remote._client = mock_client - fetch_method = getattr(remote, REMOTE_METHODS[resource_type]) - flyte_entity_latest = fetch_method(name="n1", version="latest") - flyte_entity_latest_implicit = fetch_method(name="n1") - flyte_entity_old = fetch_method(name="n1", version="old") - - assert flyte_entity_latest.entity_type_text == ENTITY_TYPE_TEXT[resource_type] - assert flyte_entity_latest.id == admin_entities[0].id - assert flyte_entity_latest.id == flyte_entity_latest_implicit.id - assert flyte_entity_latest.id != flyte_entity_old.id - - @patch("flytekit.clients.friendly.SynchronousFlyteClient") @patch("flytekit.configuration.platform.URL") @patch("flytekit.configuration.platform.INSECURE") @@ -106,39 +50,7 @@ def test_remote_fetch_workflow_execution(mock_insecure, mock_url, mock_client_ma assert flyte_workflow_execution.id == admin_workflow_execution.id -@patch("flytekit.configuration.platform.URL") -@patch("flytekit.configuration.platform.INSECURE") -def test_get_node_execution_interface(mock_insecure, mock_url): - expected_interface = TypedInterface( - {"in1": Variable(LiteralType(simple=SimpleType.STRING), "in1 description")}, - {"out1": Variable(LiteralType(simple=SimpleType.INTEGER), "out1 description")}, - ) - - node_exec_id = NodeExecutionIdentifier("node_id", WorkflowExecutionIdentifier("p1", "d1", "exec_name")) - - mock_node = MagicMock() - mock_node.id = node_exec_id.node_id - task_node = MagicMock() - flyte_task = MagicMock() - flyte_task.interface = expected_interface - task_node.flyte_task = flyte_task - mock_node.task_node = task_node - - flyte_workflow = FlyteWorkflow([mock_node], None, None, None, None, None) - - mock_url.get.return_value = "localhost" - mock_insecure.get.return_value = True - mock_client = MagicMock() - - remote = FlyteRemote.from_config("p1", "d1") - remote._client = mock_client - actual_interface = remote._get_node_execution_interface( - NodeExecution(node_exec_id, None, None, NodeExecutionMetaData(None, True, None)), flyte_workflow - ) - assert actual_interface == expected_interface - - -@patch("flytekit.remote.workflow_execution.FlyteWorkflowExecution.promote_from_model") +@patch("flytekit.remote.executions.FlyteWorkflowExecution.promote_from_model") @patch("flytekit.configuration.platform.URL") @patch("flytekit.configuration.platform.INSECURE") def test_underscore_execute_uses_launch_plan_attributes(mock_insecure, mock_url, mock_wf_exec): @@ -171,7 +83,7 @@ def local_assertions(*args, **kwargs): ) -@patch("flytekit.remote.workflow_execution.FlyteWorkflowExecution.promote_from_model") +@patch("flytekit.remote.executions.FlyteWorkflowExecution.promote_from_model") @patch("flytekit.configuration.auth.ASSUMABLE_IAM_ROLE") @patch("flytekit.configuration.platform.URL") @patch("flytekit.configuration.platform.INSECURE") @@ -201,7 +113,7 @@ def local_assertions(*args, **kwargs): ) -@patch("flytekit.remote.workflow_execution.FlyteWorkflowExecution.promote_from_model") +@patch("flytekit.remote.executions.FlyteWorkflowExecution.promote_from_model") @patch("flytekit.configuration.platform.URL") @patch("flytekit.configuration.platform.INSECURE") def test_execute_with_wrong_input_key(mock_insecure, mock_url, mock_wf_exec): diff --git a/tests/flytekit/unit/remote/test_wrapper_classes.py b/tests/flytekit/unit/remote/test_wrapper_classes.py index d37fd738c2..b26253e1db 100644 --- a/tests/flytekit/unit/remote/test_wrapper_classes.py +++ b/tests/flytekit/unit/remote/test_wrapper_classes.py @@ -68,9 +68,6 @@ def wf(b: int) -> int: assert list(fwf.interface.inputs.keys()) == ["b"] assert len(fwf.nodes) == 1 assert len(fwf.flyte_nodes) == 1 - flyte_subwfs = fwf.get_sub_workflows() - assert len(flyte_subwfs) == 1 - assert fwf.nodes[0].workflow_node.sub_workflow_ref == flyte_subwfs[0].id # Test another subwf that calls a launch plan instead of the sub_wf directly @workflow @@ -88,8 +85,6 @@ def wf2(b: int) -> int: assert list(fwf.interface.inputs.keys()) == ["b"] assert len(fwf.nodes) == 1 assert len(fwf.flyte_nodes) == 1 - flyte_subwfs = fwf.get_sub_workflows() - assert len(flyte_subwfs) == 0 # The resource type will be different, so just check the name assert fwf.nodes[0].workflow_node.launchplan_ref.name == list(lp_specs.values())[0].workflow_id.name From 9aa161bdc33299c2ec296e152d628327bf050886 Mon Sep 17 00:00:00 2001 From: Ketan Umare <16888709+kumare3@users.noreply.github.com> Date: Sun, 28 Nov 2021 13:24:22 -0800 Subject: [PATCH 024/128] Extras: Shell task (#747) Signed-off-by: maximsmol --- flytekit/extras/tasks/__init__.py | 0 flytekit/extras/tasks/shell.py | 224 ++++++++++++++++++ tests/flytekit/unit/extras/tasks/__init__.py | 0 .../flytekit/unit/extras/tasks/test_shell.py | 171 +++++++++++++ .../unit/extras/tasks/testdata/script.sh | 6 + .../unit/extras/tasks/testdata/test.csv | 0 6 files changed, 401 insertions(+) create mode 100644 flytekit/extras/tasks/__init__.py create mode 100644 flytekit/extras/tasks/shell.py create mode 100644 tests/flytekit/unit/extras/tasks/__init__.py create mode 100644 tests/flytekit/unit/extras/tasks/test_shell.py create mode 100644 tests/flytekit/unit/extras/tasks/testdata/script.sh create mode 100644 tests/flytekit/unit/extras/tasks/testdata/test.csv diff --git a/flytekit/extras/tasks/__init__.py b/flytekit/extras/tasks/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py new file mode 100644 index 0000000000..6e8dbcc21b --- /dev/null +++ b/flytekit/extras/tasks/shell.py @@ -0,0 +1,224 @@ +import datetime +import logging +import os +import re +import subprocess +import typing +from dataclasses import dataclass + +from flytekit.core.context_manager import ExecutionParameters +from flytekit.core.interface import Interface +from flytekit.core.python_function_task import PythonInstanceTask +from flytekit.core.task import TaskPlugins +from flytekit.types.directory import FlyteDirectory +from flytekit.types.file import FlyteFile + + +@dataclass +class OutputLocation: + """ + Args: + var: str The name of the output variable + var_type: typing.Type The type of output variable + location: os.PathLike The location where this output variable will be written to or a regex that accepts input + vars and generates the path. Of the form ``"{{ .inputs.v }}.tmp.md"``. + This example for a given input v, at path `/tmp/abc.csv` will resolve to `/tmp/abc.csv.tmp.md` + """ + + var: str + var_type: typing.Type + location: typing.Union[os.PathLike, str] + + +def _stringify(v: typing.Any) -> str: + """ + Special cased return for the given value. Given the type returns the string version for the type. + Handles FlyteFile and FlyteDirectory specially. Downloads and returns the downloaded filepath + """ + if isinstance(v, FlyteFile): + v.download() + return v.path + if isinstance(v, FlyteDirectory): + v.download() + return v.path + if isinstance(v, datetime.datetime): + return v.isoformat() + return str(v) + + +def _interpolate(tmpl: str, regex: re.Pattern, validate_all_match: bool = True, **kwargs) -> str: + """ + Substitutes all templates that match the supplied regex + with the given inputs and returns the substituted string. The result is non destructive towards the given string. + """ + modified = tmpl + matched = set() + for match in regex.finditer(tmpl): + expr = match.groups()[0] + var = match.groups()[1] + if var not in kwargs: + raise ValueError(f"Variable {var} in Query (part of {expr}) not found in inputs {kwargs.keys()}") + matched.add(var) + val = kwargs[var] + # str conversion should be deliberate, with right conversion for each type + modified = modified.replace(expr, _stringify(val)) + + if validate_all_match: + if len(matched) < len(kwargs.keys()): + diff = set(kwargs.keys()).difference(matched) + raise ValueError(f"Extra Inputs have no matches in script template - missing {diff}") + return modified + + +def _dummy_task_func(): + """ + A Fake function to satisfy the inner PythonTask requirements + """ + return None + + +T = typing.TypeVar("T") + + +class ShellTask(PythonInstanceTask[T]): + """ """ + + _INPUT_REGEX = re.compile(r"({{\s*.inputs.(\w+)\s*}})", re.IGNORECASE) + _OUTPUT_REGEX = re.compile(r"({{\s*.outputs.(\w+)\s*}})", re.IGNORECASE) + + def __init__( + self, + name: str, + debug: bool = False, + script: typing.Optional[str] = None, + script_file: typing.Optional[str] = None, + task_config: T = None, + inputs: typing.Optional[typing.Dict[str, typing.Type]] = None, + output_locs: typing.Optional[typing.List[OutputLocation]] = None, + **kwargs, + ): + """ + Args: + name: str Name of the Task. Should be unique in the project + debug: bool Print the generated script and other debugging information + script: The actual script specified as a string + script_file: A path to the file that contains the script (Only script or script_file) can be provided + task_config: T Configuration for the task, can be either a Pod (or coming soon, BatchJob) config + inputs: A Dictionary of input names to types + output_locs: A list of :py:class:`OutputLocations` + **kwargs: Other arguments that can be passed to :ref:class:`PythonInstanceTask` + """ + if script and script_file: + raise ValueError("Only either of script or script_file can be provided") + if not script and not script_file: + raise ValueError("Either a script or script_file is needed") + if script_file: + if not os.path.exists(script_file): + raise ValueError(f"FileNotFound: the specified Script file at path {script_file} cannot be loaded") + script_file = os.path.abspath(script_file) + + if task_config is not None: + if str(type(task_config)) != "flytekitplugins.pod.task.Pod": + raise ValueError("TaskConfig can either be empty - indicating simple container task or a PodConfig.") + + # Each instance of NotebookTask instantiates an underlying task with a dummy function that will only be used + # to run pre- and post- execute functions using the corresponding task plugin. + # We rename the function name here to ensure the generated task has a unique name and avoid duplicate task name + # errors. + # This seem like a hack. We should use a plugin_class that doesn't require a fake-function to make work. + plugin_class = TaskPlugins.find_pythontask_plugin(type(task_config)) + self._config_task_instance = plugin_class(task_config=task_config, task_function=_dummy_task_func) + # Rename the internal task so that there are no conflicts at serialization time. Technically these internal + # tasks should not be serialized at all, but we don't currently have a mechanism for skipping Flyte entities + # at serialization time. + self._config_task_instance._name = f"_bash.{name}" + self._script = script + self._script_file = script_file + self._debug = debug + self._output_locs = output_locs if output_locs else [] + outputs = self._validate_output_locs() + super().__init__( + name, + task_config, + task_type=self._config_task_instance.task_type, + interface=Interface(inputs=inputs, outputs=outputs), + **kwargs, + ) + + def _validate_output_locs(self) -> typing.Dict[str, typing.Type]: + outputs = {} + for v in self._output_locs: + if v is None: + raise ValueError("OutputLocation cannot be none") + if not isinstance(v, OutputLocation): + raise ValueError("Every output type should be an output location on the file-system") + if v.location is None: + raise ValueError(f"Output Location not provided for output var {v.var}") + if not issubclass(v.var_type, FlyteFile) and not issubclass(v.var_type, FlyteDirectory): + raise ValueError( + "Currently only outputs of type FlyteFile/FlyteDirectory and their derived types are supported" + ) + outputs[v.var] = v.var_type + return outputs + + @property + def script(self) -> typing.Optional[str]: + return self._script + + @property + def script_file(self) -> typing.Optional[os.PathLike]: + return self._script_file + + def pre_execute(self, user_params: ExecutionParameters) -> ExecutionParameters: + return self._config_task_instance.pre_execute(user_params) + + def execute(self, **kwargs) -> typing.Any: + """ + Executes the given script by substituting the inputs and outputs and extracts the outputs from the filesystem + """ + logging.info(f"Running shell script as type {self.task_type}") + if self.script_file: + with open(self.script_file) as f: + self._script = f.read() + + outputs: typing.Dict[str, str] = {} + if self._output_locs: + for v in self._output_locs: + outputs[v.var] = _interpolate(v.location, self._INPUT_REGEX, validate_all_match=False, **kwargs) + + gen_script = _interpolate(self._script, self._INPUT_REGEX, **kwargs) + # For outputs it is not necessary that all outputs are used in the script, some are implicit outputs + # for example gcc main.c will generate a.out automatically + gen_script = _interpolate(gen_script, self._OUTPUT_REGEX, validate_all_match=False, **outputs) + if self._debug: + print("\n==============================================\n") + print(gen_script) + print("\n==============================================\n") + + try: + subprocess.check_call(gen_script, shell=True) + except subprocess.CalledProcessError as e: + files = os.listdir("./") + fstr = "\n-".join(files) + logging.error( + f"Failed to Execute Script, return-code {e.returncode} \n" + f"StdErr: {e.stderr}\n" + f"StdOut: {e.stdout}\n" + f" Current directory contents: .\n-{fstr}" + ) + raise + + final_outputs = [] + for v in self._output_locs: + if issubclass(v.var_type, FlyteFile): + final_outputs.append(FlyteFile(outputs[v.var])) + if issubclass(v.var_type, FlyteDirectory): + final_outputs.append(FlyteDirectory(outputs[v.var])) + if len(final_outputs) == 1: + return final_outputs[0] + if len(final_outputs) > 1: + return tuple(final_outputs) + return None + + def post_execute(self, user_params: ExecutionParameters, rval: typing.Any) -> typing.Any: + return self._config_task_instance.post_execute(user_params, rval) diff --git a/tests/flytekit/unit/extras/tasks/__init__.py b/tests/flytekit/unit/extras/tasks/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/flytekit/unit/extras/tasks/test_shell.py b/tests/flytekit/unit/extras/tasks/test_shell.py new file mode 100644 index 0000000000..39f114a9c7 --- /dev/null +++ b/tests/flytekit/unit/extras/tasks/test_shell.py @@ -0,0 +1,171 @@ +import datetime +import os +import tempfile +from subprocess import CalledProcessError + +import pytest + +from flytekit import kwtypes +from flytekit.extras.tasks.shell import OutputLocation, ShellTask +from flytekit.types.directory import FlyteDirectory +from flytekit.types.file import CSVFile, FlyteFile + +test_file_path = os.path.dirname(os.path.realpath(__file__)) +testdata = os.path.join(test_file_path, "testdata") +script_sh = os.path.join(testdata, "script.sh") +test_csv = os.path.join(testdata, "test.csv") + + +def test_shell_task_no_io(): + t = ShellTask( + name="test", + script=""" + echo "Hello World!" + """, + ) + + t() + + +def test_shell_task_fail(): + t = ShellTask( + name="test", + script=""" + non-existent blah + """, + ) + + with pytest.raises(Exception): + t() + + +def test_input_substitution_primitive(): + t = ShellTask( + name="test", + script=""" + set -ex + cat {{ .inputs.f }} + echo "Hello World {{ .inputs.y }} on {{ .inputs.j }}" + """, + inputs=kwtypes(f=str, y=int, j=datetime.datetime), + ) + + t(f=os.path.join(test_file_path, "__init__.py"), y=5, j=datetime.datetime(2021, 11, 10, 12, 15, 0)) + t(f=os.path.join(test_file_path, "test_shell.py"), y=5, j=datetime.datetime(2021, 11, 10, 12, 15, 0)) + with pytest.raises(CalledProcessError): + t(f="non_exist.py", y=5, j=datetime.datetime(2021, 11, 10, 12, 15, 0)) + + +def test_input_substitution_files(): + t = ShellTask( + name="test", + script=""" + cat {{ .inputs.f }} + echo "Hello World {{ .inputs.y }} on {{ .inputs.j }}" + """, + inputs=kwtypes(f=CSVFile, y=FlyteDirectory, j=datetime.datetime), + ) + + assert t(f=test_csv, y=testdata, j=datetime.datetime(2021, 11, 10, 12, 15, 0)) is None + + +def test_input_output_substitution_files(): + s = """ + cat {{ .inputs.f }} > {{ .outputs.y }} + """ + t = ShellTask( + name="test", + debug=True, + script=s, + inputs=kwtypes(f=CSVFile), + output_locs=[ + OutputLocation(var="y", var_type=FlyteFile, location="{{ .inputs.f }}.mod"), + ], + ) + + assert t.script == s + + contents = "1,2,3,4\n" + with tempfile.TemporaryDirectory() as tmp: + csv = os.path.join(tmp, "abc.csv") + print(csv) + with open(csv, "w") as f: + f.write(contents) + y = t(f=csv) + assert y.path[-4:] == ".mod" + assert os.path.exists(y.path) + with open(y.path) as f: + s = f.read() + assert s == contents + + +def test_input_single_output_substitution_files(): + s = """ + cat {{ .inputs.f }} >> {{ .outputs.y }} + echo "Hello World {{ .inputs.y }} on {{ .inputs.j }}" + """ + t = ShellTask( + name="test", + debug=True, + script=s, + inputs=kwtypes(f=CSVFile, y=FlyteDirectory, j=datetime.datetime), + output_locs=[OutputLocation(var="y", var_type=FlyteFile, location="{{ .inputs.f }}.pyc")], + ) + + assert t.script == s + y = t(f=test_csv, y=testdata, j=datetime.datetime(2021, 11, 10, 12, 15, 0)) + assert y.path[-4:] == ".pyc" + + +def test_input_output_extra_var_in_template(): + t = ShellTask( + name="test", + debug=True, + script=""" + cat {{ .inputs.f }} {{ .inputs.missing }} >> {{ .outputs.y }} + echo "Hello World {{ .inputs.y }} on {{ .inputs.j }} - output {{.outputs.x}}" + """, + inputs=kwtypes(f=CSVFile, y=FlyteDirectory, j=datetime.datetime), + output_locs=[ + OutputLocation(var="x", var_type=FlyteDirectory, location="{{ .inputs.y }}"), + OutputLocation(var="y", var_type=FlyteFile, location="{{ .inputs.f }}.pyc"), + ], + ) + + with pytest.raises(ValueError): + t(f=test_csv, y=testdata, j=datetime.datetime(2021, 11, 10, 12, 15, 0)) + + +def test_input_output_extra_input(): + t = ShellTask( + name="test", + debug=True, + script=""" + cat {{ .inputs.missing }} >> {{ .outputs.y }} + echo "Hello World {{ .inputs.y }} on {{ .inputs.j }} - output {{.outputs.x}}" + """, + inputs=kwtypes(f=CSVFile, y=FlyteDirectory, j=datetime.datetime), + output_locs=[ + OutputLocation(var="x", var_type=FlyteDirectory, location="{{ .inputs.y }}"), + OutputLocation(var="y", var_type=FlyteFile, location="{{ .inputs.f }}.pyc"), + ], + ) + + with pytest.raises(ValueError): + t(f=test_csv, y=testdata, j=datetime.datetime(2021, 11, 10, 12, 15, 0)) + + +def test_shell_script(): + t = ShellTask( + name="test2", + debug=True, + script_file=script_sh, + inputs=kwtypes(f=CSVFile, y=FlyteDirectory, j=datetime.datetime), + output_locs=[ + OutputLocation(var="x", var_type=FlyteDirectory, location="{{ .inputs.y }}"), + OutputLocation(var="y", var_type=FlyteFile, location="{{ .inputs.f }}.pyc"), + ], + ) + + assert t.script_file == script_sh + t(f=test_csv, y=testdata, j=datetime.datetime(2021, 11, 10, 12, 15, 0)) diff --git a/tests/flytekit/unit/extras/tasks/testdata/script.sh b/tests/flytekit/unit/extras/tasks/testdata/script.sh new file mode 100644 index 0000000000..22012ec3ae --- /dev/null +++ b/tests/flytekit/unit/extras/tasks/testdata/script.sh @@ -0,0 +1,6 @@ +#!/bin/bash + +set -ex + +cat "{{ .inputs.f }}" >> "{{ .outputs.y }}" +echo "Hello World {{ .inputs.y }} on {{ .inputs.j }} - output {{.outputs.x}}" diff --git a/tests/flytekit/unit/extras/tasks/testdata/test.csv b/tests/flytekit/unit/extras/tasks/testdata/test.csv new file mode 100644 index 0000000000..e69de29bb2 From c430967949f9ac1fea367fad5a10fa5df4b76fc9 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Tue, 30 Nov 2021 22:42:26 +0800 Subject: [PATCH 025/128] Add support FlyteSchema in dataclass (#722) * schema in dataclass Signed-off-by: Kevin Su * Added tests Signed-off-by: Kevin Su * Fixed lint Signed-off-by: Kevin Su * Updated tests Signed-off-by: Kevin Su * Updated tests Signed-off-by: Kevin Su * Fixed lint Signed-off-by: Kevin Su * updated Signed-off-by: Kevin Su * updated Signed-off-by: Kevin Su * updated Signed-off-by: Kevin Su --- flytekit/core/type_engine.py | 65 ++++++++++++++++---- flytekit/types/schema/types.py | 20 +++--- tests/flytekit/unit/core/test_type_engine.py | 35 +++++++++++ tests/flytekit/unit/core/test_type_hints.py | 29 +++++++++ 4 files changed, 129 insertions(+), 20 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index da8cfe15f9..3c7e16e289 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -29,7 +29,7 @@ from flytekit.models import interface as _interface_models from flytekit.models import types as _type_models from flytekit.models.core import types as _core_types -from flytekit.models.literals import Literal, LiteralCollection, LiteralMap, Primitive, Scalar, Union, Void +from flytekit.models.literals import Literal, LiteralCollection, LiteralMap, Primitive, Scalar, Schema, Union, Void from flytekit.models.types import LiteralType, SimpleType, TypeStructure, UnionType try: @@ -38,15 +38,19 @@ try: from typing_extensions import get_args as _get_args except ImportError: + def _get_args(t): return t.__args__ + T = typing.TypeVar("T") DEFINITIONS = "definitions" + class TypeTransformerFailedError(TypeError, AssertionError, ValueError): ... + class TypeTransformer(typing.Generic[T]): """ Base transformer type that should be implemented for every python native type that can be handled by flytekit @@ -154,9 +158,11 @@ def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], exp def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[T]) -> T: if expected_python_type != self._type: - raise TypeTransformerFailedError(f"Cannot convert to type {expected_python_type}, only {self._type} is supported") + raise TypeTransformerFailedError( + f"Cannot convert to type {expected_python_type}, only {self._type} is supported" + ) - try: # todo(maximsmol): this is quite ugly and each transformer should really check their Literal + try: # todo(maximsmol): this is quite ugly and each transformer should really check their Literal res = self._from_literal_transformer(lv) if type(res) != self._type: raise TypeTransformerFailedError(f"Cannot convert literal {lv} to {self._type}") @@ -274,10 +280,39 @@ def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], exp raise TypeTransformerFailedError( f"Dataclass {python_type} should be decorated with @dataclass_json to be " f"serialized correctly" ) + self._serialize_flyte_type(python_val, python_type) return Literal( scalar=Scalar(generic=_json_format.Parse(cast(DataClassJsonMixin, python_val).to_json(), _struct.Struct())) ) + def _serialize_flyte_type(self, python_val: T, python_type: Type[T]): + """ + If any field inside the dataclass is flyte type, we should use flyte type transformer for that field. + """ + from flytekit.types.schema.types import FlyteSchema, FlyteSchemaTransformer + + for f in dataclasses.fields(python_type): + v = python_val.__getattribute__(f.name) + if inspect.isclass(f.type) and issubclass(f.type, FlyteSchema): + FlyteSchemaTransformer().to_literal(FlyteContext.current_context(), v, f.type, None) + elif dataclasses.is_dataclass(f.type): + self._serialize_flyte_type(v, f.type) + + def _deserialize_flyte_type(self, python_val: T, expected_python_type: Type["FlyteSchema"]): + from flytekit.types.schema.types import FlyteSchema, FlyteSchemaTransformer + + for f in dataclasses.fields(expected_python_type): + v = python_val.__getattribute__(f.name) + if inspect.isclass(f.type) and issubclass(f.type, FlyteSchema): + t = FlyteSchemaTransformer() + t.to_python_value( + FlyteContext.current_context(), + Literal(scalar=Scalar(schema=Schema(v.remote_path, t._get_schema_type(f.type)))), + f.type, + ) + elif dataclasses.is_dataclass(f.type): + self._deserialize_flyte_type(v, f.type) + def _fix_val_int(self, t: typing.Type, val: typing.Any) -> typing.Any: if t == int: return int(val) @@ -320,7 +355,9 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: f"Dataclass {expected_python_type} should be decorated with @dataclass_json to be " f"serialized correctly" ) + dc = cast(DataClassJsonMixin, expected_python_type).from_json(_json_format.MessageToJson(lv.scalar.generic)) + self._deserialize_flyte_type(dc, expected_python_type) return self._fix_dataclass_int(expected_python_type, dc) def guess_python_type(self, literal_type: LiteralType) -> Type[T]: @@ -633,10 +670,12 @@ def guess_python_type(self, literal_type: LiteralType) -> Type[list]: return typing.List[ct] raise ValueError(f"List transformer cannot reverse {literal_type}") + def _add_tag_to_type(x: LiteralType, tag: str) -> LiteralType: x._structure = TypeStructure(tag=tag) return x + class UnionTransformer(TypeTransformer[T]): """ Transformer that handles a typing.Union[T1, T2, ...] @@ -678,7 +717,7 @@ def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], exp def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[T]) -> Optional[typing.Any]: union_tag = None - if lv.scalar is not None and lv.scalar.union is not None: + if lv.scalar is not None and lv.scalar.union is not None: union_type = lv.scalar.union.stored_type if union_type.structure is not None: union_tag = union_type.structure.tag @@ -693,22 +732,24 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: if trans.name != union_tag: continue - assert lv.scalar is not None # type checker - assert lv.scalar.union is not None # type checker + assert lv.scalar is not None # type checker + assert lv.scalar.union is not None # type checker res = trans.to_python_value(ctx, lv.scalar.union.value, v) res_tag = trans.name if found_res: raise TypeError( - "Ambiguous choice of variant for union type. " + - f"Both {res_tag} and {trans.name} transformers match") + "Ambiguous choice of variant for union type. " + + f"Both {res_tag} and {trans.name} transformers match" + ) found_res = True else: res = trans.to_python_value(ctx, lv, v) if found_res: raise TypeError( - "Ambiguous choice of variant for union type. " + - f"Both {res_tag} and {trans.name} transformers match") + "Ambiguous choice of variant for union type. " + + f"Both {res_tag} and {trans.name} transformers match" + ) res_tag = trans.name found_res = True except TypeTransformerFailedError as e: @@ -976,11 +1017,13 @@ def _check_and_covert_float(lv: Literal) -> float: return float(lv.scalar.primitive.integer) raise TypeTransformerFailedError(f"Cannot convert literal {lv} to float") + def _check_and_convert_void(lv: Literal) -> None: if lv.scalar.none_type is None: raise TypeTransformerFailedError(f"Cannot conver literal {lv} to None") return None + def _register_default_type_transformers(): TypeEngine.register( SimpleTransformer( @@ -1048,7 +1091,7 @@ def _register_default_type_transformers(): type(None), _type_models.LiteralType(simple=_type_models.SimpleType.NONE), lambda x: Literal(scalar=Scalar(none_type=Void())), - lambda x: _check_and_convert_void(x) + lambda x: _check_and_convert_void(x), ), [None], ) diff --git a/flytekit/types/schema/types.py b/flytekit/types/schema/types.py index 9864f660b4..2cd2450393 100644 --- a/flytekit/types/schema/types.py +++ b/flytekit/types/schema/types.py @@ -4,11 +4,13 @@ import os import typing from abc import abstractmethod -from dataclasses import dataclass +from dataclasses import dataclass, field from enum import Enum from typing import Type import numpy as _np +from dataclasses_json import config, dataclass_json +from marshmallow import fields from flytekit.core.context_manager import FlyteContext, FlyteContextManager from flytekit.core.type_engine import T, TypeEngine, TypeTransformer, TypeTransformerFailedError @@ -167,7 +169,10 @@ def get_handler(cls, t: Type) -> SchemaHandler: return cls._SCHEMA_HANDLERS[t] +@dataclass_json +@dataclass class FlyteSchema(object): + remote_path: typing.Optional[os.PathLike] = field(default=None, metadata=config(mm_field=fields.String())) """ This is the main schema class that users should use. """ @@ -220,7 +225,7 @@ def format(cls) -> SchemaFormat: def __init__( self, local_path: os.PathLike = None, - remote_path: str = None, + remote_path: os.PathLike = None, supported_mode: SchemaOpenMode = SchemaOpenMode.WRITE, downloader: typing.Callable[[str, os.PathLike], None] = None, ): @@ -234,10 +239,11 @@ def __init__( ): raise ValueError("To create a FlyteSchema in write mode, local_path is required") - if local_path is None: - local_path = FlyteContextManager.current_context().file_access.get_random_local_directory() + local_path = local_path or FlyteContextManager.current_context().file_access.get_random_local_directory() self._local_path = local_path - self._remote_path = remote_path + # Make this field public, so that the dataclass transformer can set a value for it + # https://github.com/flyteorg/flytekit/blob/bcc8541bd6227b532f8462563fe8aac902242b21/flytekit/core/type_engine.py#L298 + self.remote_path = remote_path or FlyteContextManager.current_context().file_access.get_random_remote_path() self._supported_mode = supported_mode # This is a special attribute that indicates if the data was either downloaded or uploaded self._downloaded = False @@ -247,10 +253,6 @@ def __init__( def local_path(self) -> os.PathLike: return self._local_path - @property - def remote_path(self) -> str: - return typing.cast(str, self._remote_path) - @property def supported_mode(self) -> SchemaOpenMode: return self._supported_mode diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 1958020065..2ffcc781b2 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -5,6 +5,7 @@ from datetime import timedelta from enum import Enum +import pandas as pd import pytest from dataclasses_json import DataClassJsonMixin, dataclass_json from flyteidl.core import errors_pb2 @@ -14,6 +15,8 @@ from marshmallow_jsonschema import JSONSchema import flytekit.common.exceptions.user as user_exceptions +from flytekit import kwtypes +from flytekit.common.exceptions import user as user_exceptions from flytekit.common.types import primitives from flytekit.core.context_manager import FlyteContext, FlyteContextManager from flytekit.core.type_engine import ( @@ -36,6 +39,7 @@ from flytekit.types.file.file import FlyteFile, FlyteFilePathTransformer from flytekit.types.pickle import FlytePickle from flytekit.types.pickle.pickle import FlytePickleTransformer +from flytekit.types.schema import FlyteSchema T = typing.TypeVar("T") @@ -882,3 +886,34 @@ def test_dict_to_literal_map_with_wrong_input_type(): guessed_python_types = {"a": str} with pytest.raises(user_exceptions.FlyteTypeException): TypeEngine.dict_to_literal_map(ctx, input, guessed_python_types) + + +TestSchema = FlyteSchema[kwtypes(some_str=str)] + + +@dataclass_json +@dataclass +class InnerResult: + number: int + schema: TestSchema + + +@dataclass_json +@dataclass +class Result: + result: InnerResult + schema: TestSchema + + +def test_schema_in_dataclass(): + schema = TestSchema() + df = pd.DataFrame(data={"some_str": ["a", "b", "c"]}) + schema.open().write(df) + o = Result(result=InnerResult(number=1, schema=schema), schema=schema) + ctx = FlyteContext.current_context() + tf = DataclassTransformer() + lt = tf.get_literal_type(Result) + lv = tf.to_literal(ctx, o, Result, lt) + ot = tf.to_python_value(ctx, lv=lv, expected_python_type=Result) + + assert o == ot diff --git a/tests/flytekit/unit/core/test_type_hints.py b/tests/flytekit/unit/core/test_type_hints.py index 05a8319a82..57de0c29b2 100644 --- a/tests/flytekit/unit/core/test_type_hints.py +++ b/tests/flytekit/unit/core/test_type_hints.py @@ -1091,6 +1091,35 @@ def wf(x: int) -> Datum: assert wf(x=10) == Datum(10, Color.RED) +def test_flyte_schema_dataclass(): + TestSchema = FlyteSchema[kwtypes(some_str=str)] + + @dataclass_json + @dataclass + class InnerResult: + number: int + schema: TestSchema + + @dataclass_json + @dataclass + class Result: + result: InnerResult + schema: TestSchema + + schema = TestSchema() + + @task + def t1(x: int) -> Result: + + return Result(result=InnerResult(number=x, schema=schema), schema=schema) + + @workflow + def wf(x: int) -> Result: + return t1(x=x) + + assert wf(x=10) == Result(result=InnerResult(number=10, schema=schema), schema=schema) + + def test_environment(): @task(environment={"FOO": "foofoo", "BAZ": "baz"}) def t1(a: int) -> str: From fac636d6bbd5fb254b8828b6ecf8272859b1a781 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Thu, 2 Dec 2021 02:54:42 +0800 Subject: [PATCH 026/128] Remove workflow_execution.py (#758) Signed-off-by: Kevin Su Signed-off-by: maximsmol --- flytekit/clis/sdk_in_container/placeholders.py | 0 flytekit/common/types/files.py | 0 flytekit/remote/workflow_execution.py | 0 tests/flytekit/unit/configuration/configs/empty.config | 0 4 files changed, 0 insertions(+), 0 deletions(-) delete mode 100644 flytekit/clis/sdk_in_container/placeholders.py delete mode 100644 flytekit/common/types/files.py delete mode 100644 flytekit/remote/workflow_execution.py delete mode 100644 tests/flytekit/unit/configuration/configs/empty.config diff --git a/flytekit/clis/sdk_in_container/placeholders.py b/flytekit/clis/sdk_in_container/placeholders.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/flytekit/common/types/files.py b/flytekit/common/types/files.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/flytekit/remote/workflow_execution.py b/flytekit/remote/workflow_execution.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/tests/flytekit/unit/configuration/configs/empty.config b/tests/flytekit/unit/configuration/configs/empty.config deleted file mode 100644 index e69de29bb2..0000000000 From 4036507b8d13468763a41b8b9c84063811398dd1 Mon Sep 17 00:00:00 2001 From: Ketan Umare <16888709+kumare3@users.noreply.github.com> Date: Thu, 2 Dec 2021 10:08:28 -0800 Subject: [PATCH 027/128] Get raw input/output from remote execution (#675) * [wip] for feast demo Signed-off-by: Ketan Umare * clean up a bit Signed-off-by: Yee Hing Tong * add a test and move where constructor is called Signed-off-by: Yee Hing Tong * remove unneeded import Signed-off-by: Yee Hing Tong * add a part of a test Signed-off-by: Yee Hing Tong * Added tests Signed-off-by: Kevin Su * Fixed lint Signed-off-by: Kevin Su * typo Signed-off-by: Kevin Su Co-authored-by: Yee Hing Tong Co-authored-by: Kevin Su Signed-off-by: maximsmol --- flytekit/core/type_engine.py | 23 ++++++++ flytekit/remote/executions.py | 15 +++++ flytekit/remote/remote.py | 10 +++- .../integration/remote/test_remote.py | 2 + tests/flytekit/unit/core/test_type_engine.py | 59 +++++++++++++++++++ 5 files changed, 106 insertions(+), 3 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 3c7e16e289..2550d6f73f 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -1114,4 +1114,27 @@ def _register_default_type_transformers(): TypeEngine.register_restricted_type("named tuple", NamedTuple) +class LiteralsResolver(object): + """ + LiteralsResolver is a helper class meant primarily for use with the FlyteRemote experience or any other situation + where you might be working with LiteralMaps. This object allows the caller to specify the Python type that should + correspond to an element of the map. + TODO: Add an optional Flyte idl interface model object to the constructor + """ + + def __init__(self, literals: typing.Dict[str, Literal]): + self._literals = literals + + @property + def literals(self): + return self._literals + + def get(self, attr: str, as_type: Optional[typing.Type] = None): + if attr not in self._literals: + raise AttributeError(f"Attribute {attr} not found") + if as_type is None: + raise ValueError("as_type argument can't be None yet.") + return TypeEngine.to_python_value(FlyteContext.current_context(), self._literals[attr], as_type) + + _register_default_type_transformers() diff --git a/flytekit/remote/executions.py b/flytekit/remote/executions.py index 05b94b2302..b581769f52 100644 --- a/flytekit/remote/executions.py +++ b/flytekit/remote/executions.py @@ -4,6 +4,7 @@ from flytekit.common.exceptions import user as _user_exceptions from flytekit.common.exceptions import user as user_exceptions +from flytekit.core.type_engine import LiteralsResolver from flytekit.models import execution as execution_models from flytekit.models import node_execution as node_execution_models from flytekit.models.admin import task_execution as admin_task_execution_models @@ -83,6 +84,8 @@ def __init__(self, *args, **kwargs): self._inputs = None self._outputs = None self._flyte_workflow: Optional[FlyteWorkflow] = None + self._raw_inputs: Optional[LiteralsResolver] = None + self._raw_outputs: Optional[LiteralsResolver] = None @property def node_executions(self) -> Dict[str, "FlyteNodeExecution"]: @@ -111,6 +114,18 @@ def outputs(self) -> Dict[str, Any]: raise _user_exceptions.FlyteAssertion("Outputs could not be found because the execution ended in failure.") return self._outputs + @property + def raw_outputs(self) -> LiteralsResolver: + if self._raw_outputs is None: + raise ValueError(f"WF execution: {self} doesn't have raw outputs set") + return self._raw_outputs + + @property + def raw_inputs(self) -> LiteralsResolver: + if self._raw_inputs is None: + raise ValueError(f"WF execution: {self} doesn't have raw inputs set") + return self._raw_inputs + @property def error(self) -> core_execution_models.ExecutionError: """ diff --git a/flytekit/remote/remote.py b/flytekit/remote/remote.py index 04e18d6794..117abf27e3 100644 --- a/flytekit/remote/remote.py +++ b/flytekit/remote/remote.py @@ -48,7 +48,7 @@ from flytekit.core.context_manager import FlyteContextManager, ImageConfig, SerializationSettings, get_image_config from flytekit.core.data_persistence import FileAccessProvider from flytekit.core.launch_plan import LaunchPlan -from flytekit.core.type_engine import TypeEngine +from flytekit.core.type_engine import LiteralsResolver, TypeEngine from flytekit.core.workflow import WorkflowBase from flytekit.models import common as common_models from flytekit.models import launch_plan as launch_plan_models @@ -1255,15 +1255,19 @@ def _assign_inputs_and_outputs( ): """Helper for assigning synced inputs and outputs to an execution object.""" with self.remote_context() as ctx: + input_literal_map = self._get_input_literal_map(execution_data) + execution._raw_inputs = LiteralsResolver(input_literal_map.literals) execution._inputs = TypeEngine.literal_map_to_kwargs( ctx=ctx, - lm=self._get_input_literal_map(execution_data), + lm=input_literal_map, python_types=TypeEngine.guess_python_types(interface.inputs), ) if execution.is_complete and not execution.error: + output_literal_map = self._get_output_literal_map(execution_data) + execution._raw_outputs = LiteralsResolver(output_literal_map.literals) execution._outputs = TypeEngine.literal_map_to_kwargs( ctx=ctx, - lm=self._get_output_literal_map(execution_data), + lm=output_literal_map, python_types=TypeEngine.guess_python_types(interface.outputs), ) return execution diff --git a/tests/flytekit/integration/remote/test_remote.py b/tests/flytekit/integration/remote/test_remote.py index 9dc32a2e54..a45c037279 100644 --- a/tests/flytekit/integration/remote/test_remote.py +++ b/tests/flytekit/integration/remote/test_remote.py @@ -156,6 +156,8 @@ def test_fetch_execute_task(flyteclient, flyte_workflows_register): execution = remote.execute(flyte_task, {"a": 10}, wait=True) assert execution.outputs["t1_int_output"] == 12 assert execution.outputs["c"] == "world" + assert execution.raw_inputs.get("a", int) == 10 + assert execution.raw_outputs.get("c", str) == "world" def test_execute_python_task(flyteclient, flyte_workflows_register, flyte_remote_env): diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 2ffcc781b2..8463a4b6cd 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -23,6 +23,7 @@ DataclassTransformer, DictTransformer, ListTransformer, + LiteralsResolver, SimpleTransformer, TypeEngine, TypeTransformer, @@ -917,3 +918,61 @@ def test_schema_in_dataclass(): ot = tf.to_python_value(ctx, lv=lv, expected_python_type=Result) assert o == ot + + +@pytest.mark.parametrize( + "literal_value,python_type,expected_python_value", + [ + ( + Literal( + collection=LiteralCollection( + literals=[ + Literal(scalar=Scalar(primitive=Primitive(integer=1))), + Literal(scalar=Scalar(primitive=Primitive(integer=2))), + Literal(scalar=Scalar(primitive=Primitive(integer=3))), + ] + ) + ), + typing.List[int], + [1, 2, 3], + ), + ( + Literal( + map=LiteralMap( + literals={ + "k1": Literal(scalar=Scalar(primitive=Primitive(string_value="v1"))), + "k2": Literal(scalar=Scalar(primitive=Primitive(string_value="2"))), + }, + ) + ), + typing.Dict[str, str], + {"k1": "v1", "k2": "2"}, + ), + ], +) +def test_literals_resolver(literal_value, python_type, expected_python_value): + lit_dict = {"a": literal_value} + + lr = LiteralsResolver(lit_dict) + out = lr.get("a", python_type) + assert out == expected_python_value + + +def test_guess_of_dataclass(): + @dataclass_json + @dataclass() + class Foo(object): + x: int + y: str + z: typing.Dict[str, int] + + def hello(self): + ... + + lt = TypeEngine.to_literal_type(Foo) + foo = Foo(1, "hello", {"world": 3}) + lv = TypeEngine.to_literal(FlyteContext.current_context(), foo, Foo, lt) + lit_dict = {"a": lv} + lr = LiteralsResolver(lit_dict) + assert lr.get("a", Foo) == foo + assert hasattr(lr.get("a", Foo), "hello") is True From 567d8c1326eee5da453a243267eed18c491a373f Mon Sep 17 00:00:00 2001 From: Lisa <30621230+aeioulisa@users.noreply.github.com> Date: Tue, 7 Dec 2021 01:51:09 +0800 Subject: [PATCH 028/128] Fix mypy errors in flytekit/types (#757) Signed-off-by: Lisa Signed-off-by: Kevin Su --- flytekit/types/directory/__init__.py | 3 ++- flytekit/types/directory/types.py | 6 +++--- flytekit/types/file/__init__.py | 30 ++++++++++++++++++--------- flytekit/types/file/file.py | 10 ++++----- flytekit/types/pickle/pickle.py | 6 +++--- flytekit/types/schema/types.py | 27 ++++++++++++++---------- flytekit/types/schema/types_pandas.py | 4 ++-- 7 files changed, 51 insertions(+), 35 deletions(-) diff --git a/flytekit/types/directory/__init__.py b/flytekit/types/directory/__init__.py index 4edb5f9205..97d8ab57ce 100644 --- a/flytekit/types/directory/__init__.py +++ b/flytekit/types/directory/__init__.py @@ -18,7 +18,8 @@ # The following section provides some predefined aliases for commonly used FlyteDirectory formats. -TensorboardLogs = FlyteDirectory[typing.TypeVar("tensorboard")] +tensorboard = typing.TypeVar("tensorboard") +TensorboardLogs = FlyteDirectory[tensorboard] """ This type can be used to denote that the output is a folder that contains logs that can be loaded in tensorboard. this is usually the SummaryWriter output in pytorch or Keras callbacks which record the history readable by diff --git a/flytekit/types/directory/types.py b/flytekit/types/directory/types.py index a0112c1eea..81f4fb0fd0 100644 --- a/flytekit/types/directory/types.py +++ b/flytekit/types/directory/types.py @@ -133,7 +133,7 @@ def __fspath__(self): def extension(cls) -> str: return "" - def __class_getitem__(cls, item: typing.Type) -> typing.Type[FlyteDirectory]: + def __class_getitem__(cls, item: typing.Union[typing.Type, str]) -> typing.Type[FlyteDirectory]: if item is None: return cls item_string = str(item) @@ -290,7 +290,7 @@ def _downloader(): expected_format = self.get_format(expected_python_type) - fd = FlyteDirectory[expected_format](local_folder, _downloader) + fd = FlyteDirectory.__class_getitem__(expected_format)(local_folder, _downloader) fd._remote_source = uri return fd @@ -300,7 +300,7 @@ def guess_python_type(self, literal_type: LiteralType) -> typing.Type[FlyteDirec literal_type.blob is not None and literal_type.blob.dimensionality == _core_types.BlobType.BlobDimensionality.MULTIPART ): - return FlyteDirectory[typing.TypeVar(literal_type.blob.format)] + return FlyteDirectory.__class_getitem__(literal_type.blob.format) raise ValueError(f"Transformer {self} cannot reverse {literal_type}") diff --git a/flytekit/types/file/__init__.py b/flytekit/types/file/__init__.py index 2b65efbcd6..81796fc49e 100644 --- a/flytekit/types/file/__init__.py +++ b/flytekit/types/file/__init__.py @@ -28,59 +28,69 @@ # This makes their usage extremely simple for the users. Please keep the list sorted. -HDF5EncodedFile = FlyteFile[typing.TypeVar("hdf5")] +hdf5 = typing.TypeVar("hdf5") +HDF5EncodedFile = FlyteFile[hdf5] """ This can be used to denote that the returned file is of type hdf5 and can be received by other tasks that accept an hdf5 format. This is usually useful for serializing Tensorflow models """ -HTMLPage = FlyteFile[typing.TypeVar("html")] +html = typing.TypeVar("html") +HTMLPage = FlyteFile[html] """ Can be used to receive or return an PNGImage. The underlying type is a FlyteFile, type. This is just a decoration and useful for attaching content type information with the file and automatically documenting code. """ -JoblibSerializedFile = FlyteFile[typing.TypeVar("joblib")] +joblib = typing.TypeVar("joblib") +JoblibSerializedFile = FlyteFile[joblib] """ This File represents a file that was serialized using `joblib.dump` method can be loaded back using `joblib.load` """ -JPEGImageFile = FlyteFile[typing.TypeVar("jpeg")] +jpeg = typing.TypeVar("jpeg") +JPEGImageFile = FlyteFile[jpeg] """ Can be used to receive or return an JPEGImage. The underlying type is a FlyteFile, type. This is just a decoration and useful for attaching content type information with the file and automatically documenting code. """ -PDFFile = FlyteFile[typing.TypeVar("pdf")] +pdf = typing.TypeVar("pdf") +PDFFile = FlyteFile[pdf] """ Can be used to receive or return an PDFFile. The underlying type is a FlyteFile, type. This is just a decoration and useful for attaching content type information with the file and automatically documenting code. """ -PNGImageFile = FlyteFile[typing.TypeVar("png")] +png = typing.TypeVar("png") +PNGImageFile = FlyteFile[png] """ Can be used to receive or return an PNGImage. The underlying type is a FlyteFile, type. This is just a decoration and useful for attaching content type information with the file and automatically documenting code. """ -PythonPickledFile = FlyteFile[typing.TypeVar("python-pickle")] +python_pickle = typing.TypeVar("python_pickle") +PythonPickledFile = FlyteFile[python_pickle] """ This type can be used when a serialized python pickled object is returned and shared between tasks. This only adds metadata to the file in Flyte, but does not really carry any object information """ -PythonNotebook = FlyteFile[typing.TypeVar("ipynb")] +ipynb = typing.TypeVar("ipynb") +PythonNotebook = FlyteFile[ipynb] """ This type is used to identify a python notebook file """ -SVGImageFile = FlyteFile[typing.TypeVar("svg")] +svg = typing.TypeVar("svg") +SVGImageFile = FlyteFile[svg] """ Can be used to receive or return an SVGImage. The underlying type is a FlyteFile, type. This is just a decoration and useful for attaching content type information with the file and automatically documenting code. """ -CSVFile = FlyteFile[typing.TypeVar("csv")] +csv = typing.TypeVar("csv") +CSVFile = FlyteFile[csv] """ Can be used to receive or return a CSVFile. The underlying type is a FlyteFile, type. This is just a decoration and useful for attaching content type information with the file and automatically documenting code. diff --git a/flytekit/types/file/file.py b/flytekit/types/file/file.py index 2421823d6c..4887976b51 100644 --- a/flytekit/types/file/file.py +++ b/flytekit/types/file/file.py @@ -137,7 +137,7 @@ def t2() -> flytekit_typing.FlyteFile["csv"]: def extension(cls) -> str: return "" - def __class_getitem__(cls, item: typing.Type) -> typing.Type[FlyteFile]: + def __class_getitem__(cls, item: typing.Union[str, typing.Type]) -> typing.Type[FlyteFile]: if item is None: return cls item_string = str(item) @@ -220,10 +220,10 @@ def __init__(self): super().__init__(name="FlyteFilePath", t=FlyteFile) @staticmethod - def get_format(t: typing.Union[typing.Type[FlyteFile]]) -> str: + def get_format(t: typing.Union[typing.Type[FlyteFile], os.PathLike]) -> str: if t is os.PathLike: return "" - return t.extension() + return typing.cast(FlyteFile, t).extension() def _blob_type(self, format: str) -> BlobType: return BlobType(format=format, dimensionality=BlobType.BlobDimensionality.SINGLE) @@ -342,14 +342,14 @@ def _downloader(): return ctx.file_access.get_data(uri, local_path, is_multipart=False) expected_format = FlyteFilePathTransformer.get_format(expected_python_type) - ff = FlyteFile[expected_format](local_path, _downloader) + ff = FlyteFile.__class_getitem__(expected_format)(local_path, _downloader) ff._remote_source = uri return ff def guess_python_type(self, literal_type: LiteralType) -> typing.Type[FlyteFile[typing.Any]]: if literal_type.blob is not None and literal_type.blob.dimensionality == BlobType.BlobDimensionality.SINGLE: - return FlyteFile[typing.TypeVar(literal_type.blob.format)] + return FlyteFile.__class_getitem__(literal_type.blob.format) raise ValueError(f"Transformer {self} cannot reverse {literal_type}") diff --git a/flytekit/types/pickle/pickle.py b/flytekit/types/pickle/pickle.py index 8251d111bb..9219d3a8b4 100644 --- a/flytekit/types/pickle/pickle.py +++ b/flytekit/types/pickle/pickle.py @@ -20,10 +20,10 @@ class FlytePickle(typing.Generic[T]): """ @classmethod - def python_type(cls) -> None: - return None + def python_type(cls) -> typing.Type: + return type(None) - def __class_getitem__(cls, python_type: typing.Type) -> typing.Type[T]: + def __class_getitem__(cls, python_type: typing.Type) -> typing.Type: if python_type is None: return cls diff --git a/flytekit/types/schema/types.py b/flytekit/types/schema/types.py index 2cd2450393..6b44af5396 100644 --- a/flytekit/types/schema/types.py +++ b/flytekit/types/schema/types.py @@ -6,6 +6,7 @@ from abc import abstractmethod from dataclasses import dataclass, field from enum import Enum +from pathlib import Path from typing import Type import numpy as _np @@ -18,6 +19,8 @@ from flytekit.models.types import LiteralType, SchemaType from flytekit.plugins import pandas +T = typing.TypeVar("T") + class SchemaFormat(Enum): """ @@ -37,7 +40,7 @@ class SchemaOpenMode(Enum): WRITE = "w" -def generate_ordered_files(directory: os.PathLike, n: int) -> str: +def generate_ordered_files(directory: os.PathLike, n: int) -> typing.Generator[str, None, None]: for i in range(n): yield os.path.join(directory, f"{i:05}") @@ -73,12 +76,12 @@ def all(self, **kwargs) -> T: class SchemaWriter(typing.Generic[T]): - def __init__(self, to_path: str, cols: typing.Dict[str, type], fmt: SchemaFormat): + def __init__(self, to_path: str, cols: typing.Optional[typing.Dict[str, type]], fmt: SchemaFormat): self._to_path = to_path self._fmt = fmt self._columns = cols # TODO This should be change to send a stop instead of hardcoded to 1024 - self._file_name_gen = generate_ordered_files(self._to_path, 1024) + self._file_name_gen = generate_ordered_files(Path(self._to_path), 1024) @property def to_path(self) -> str: @@ -107,14 +110,14 @@ def iter(self, **kwargs) -> typing.Generator[T, None, None]: with os.scandir(self._from_path) as it: for entry in it: if not entry.name.startswith(".") and entry.is_file(): - yield self._read(entry.path, **kwargs) + yield self._read(Path(entry.path), **kwargs) def all(self, **kwargs) -> T: - files = [] + files: typing.List[os.PathLike] = [] with os.scandir(self._from_path) as it: for entry in it: if not entry.name.startswith(".") and entry.is_file(): - files.append(entry.path) + files.append(Path(entry.path)) return self._read(*files, **kwargs) @@ -279,13 +282,15 @@ def open( if not h.handles_remote_io: # The Schema Handler does not manage its own IO, and this it will expect the files are on local file-system if self._supported_mode == SchemaOpenMode.READ and not self._downloaded: + if self._downloader is None: + raise AssertionError("downloader cannot be None in read mode!") # Only for readable objects if they are not downloaded already, we should download them # Write objects should already have everything written to self._downloader(self.remote_path, self.local_path) self._downloaded = True if mode == SchemaOpenMode.WRITE: - return h.writer(self.local_path, self.columns(), self.format()) - return h.reader(self.local_path, self.columns(), self.format()) + return h.writer(typing.cast(str, self.local_path), self.columns(), self.format()) + return h.reader(typing.cast(str, self.local_path), self.columns(), self.format()) # Remote IO is handled. So we will just pass the remote reference to the object if mode == SchemaOpenMode.WRITE: @@ -387,10 +392,10 @@ def downloader(x, y): supported_mode=SchemaOpenMode.READ, ) - def guess_python_type(self, literal_type: LiteralType) -> Type[T]: + def guess_python_type(self, literal_type: LiteralType) -> Type[FlyteSchema]: if not literal_type.schema: raise ValueError(f"Cannot reverse {literal_type}") - columns: dict[Type] = {} + columns: typing.Dict[str, Type] = {} for literal_column in literal_type.schema.columns: if literal_column.type == SchemaType.SchemaColumn.SchemaColumnType.INTEGER: columns[literal_column.name] = int @@ -406,7 +411,7 @@ def guess_python_type(self, literal_type: LiteralType) -> Type[T]: columns[literal_column.name] = bool else: raise ValueError(f"Unknown schema column type {literal_column}") - return FlyteSchema[columns] + return FlyteSchema.__class_getitem__(columns) TypeEngine.register(FlyteSchemaTransformer()) diff --git a/flytekit/types/schema/types_pandas.py b/flytekit/types/schema/types_pandas.py index 41a5423c08..0edf024b08 100644 --- a/flytekit/types/schema/types_pandas.py +++ b/flytekit/types/schema/types_pandas.py @@ -15,7 +15,7 @@ class ParquetIO(object): PARQUET_ENGINE = "pyarrow" - def _read(self, chunk: os.PathLike, columns: typing.List[str], **kwargs) -> pandas.DataFrame: + def _read(self, chunk: os.PathLike, columns: typing.Optional[typing.List[str]], **kwargs) -> pandas.DataFrame: return pandas.read_parquet(chunk, columns=columns, engine=self.PARQUET_ENGINE, **kwargs) def read(self, *files: os.PathLike, columns: typing.List[str] = None, **kwargs) -> pandas.DataFrame: @@ -59,7 +59,7 @@ def write( class FastParquetIO(ParquetIO): PARQUET_ENGINE = "fastparquet" - def _read(self, chunk: os.PathLike, columns: typing.List[str], **kwargs) -> pandas.DataFrame: + def _read(self, chunk: os.PathLike, columns: typing.Optional[typing.List[str]], **kwargs) -> pandas.DataFrame: from fastparquet import ParquetFile as _ParquetFile from fastparquet import thrift_structures as _ts From 6b32abc216e27e41bd2dc06755a0360e09cd4fc4 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Tue, 7 Dec 2021 12:32:51 +0800 Subject: [PATCH 029/128] Remote client failed to fetch FlytePickle object (#764) * Fetch pickle value from flytekit remote Signed-off-by: Kevin Su * Fix tests Signed-off-by: Kevin Su * Remove default value Signed-off-by: Kevin Su --- flytekit/remote/remote.py | 12 +++++++-- flytekit/types/file/file.py | 7 ++++- flytekit/types/pickle/pickle.py | 10 ++++++++ tests/flytekit/unit/core/test_type_engine.py | 27 ++++++++++++++++++++ 4 files changed, 53 insertions(+), 3 deletions(-) diff --git a/flytekit/remote/remote.py b/flytekit/remote/remote.py index 117abf27e3..60ae7bbf3e 100644 --- a/flytekit/remote/remote.py +++ b/flytekit/remote/remote.py @@ -235,9 +235,17 @@ def __init__( # Not exposing this as a property for now. self._entrypoint_settings = entrypoint_settings + raw_output_data_prefix = auth_config.RAW_OUTPUT_DATA_PREFIX.get() or os.path.join( + sdk_config.LOCAL_SANDBOX.get(), "control_plane_raw" + ) + self._file_access = file_access or FileAccessProvider( + local_sandbox_dir=os.path.join(sdk_config.LOCAL_SANDBOX.get(), "control_plane_metadata"), + raw_output_prefix=raw_output_data_prefix, + ) # Save the file access object locally, but also make it available for use from the context. - FlyteContextManager.with_context(FlyteContextManager.current_context().with_file_access(file_access).build()) - self._file_access = file_access + FlyteContextManager.with_context( + FlyteContextManager.current_context().with_file_access(self._file_access).build() + ) # TODO: Reconsider whether we want this. Probably best to not cache. self._serialized_entity_cache = OrderedDict() diff --git a/flytekit/types/file/file.py b/flytekit/types/file/file.py index 4887976b51..826bc3916d 100644 --- a/flytekit/types/file/file.py +++ b/flytekit/types/file/file.py @@ -10,6 +10,7 @@ from flytekit.models.core.types import BlobType from flytekit.models.literals import Blob, BlobMetadata, Literal, Scalar from flytekit.models.types import LiteralType +from flytekit.types.pickle.pickle import FlytePickleTransformer def noop(): @@ -348,7 +349,11 @@ def _downloader(): return ff def guess_python_type(self, literal_type: LiteralType) -> typing.Type[FlyteFile[typing.Any]]: - if literal_type.blob is not None and literal_type.blob.dimensionality == BlobType.BlobDimensionality.SINGLE: + if ( + literal_type.blob is not None + and literal_type.blob.dimensionality == BlobType.BlobDimensionality.SINGLE + and literal_type.blob.format != FlytePickleTransformer.PYTHON_PICKLE_FORMAT + ): return FlyteFile.__class_getitem__(literal_type.blob.format) raise ValueError(f"Transformer {self} cannot reverse {literal_type}") diff --git a/flytekit/types/pickle/pickle.py b/flytekit/types/pickle/pickle.py index 9219d3a8b4..3472dec7e6 100644 --- a/flytekit/types/pickle/pickle.py +++ b/flytekit/types/pickle/pickle.py @@ -78,6 +78,16 @@ def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], exp ctx.file_access.put_data(uri, remote_path, is_multipart=False) return Literal(scalar=Scalar(blob=Blob(metadata=meta, uri=remote_path))) + def guess_python_type(self, literal_type: LiteralType) -> typing.Type[FlytePickle[typing.Any]]: + if ( + literal_type.blob is not None + and literal_type.blob.dimensionality == _core_types.BlobType.BlobDimensionality.SINGLE + and literal_type.blob.format == FlytePickleTransformer.PYTHON_PICKLE_FORMAT + ): + return FlytePickle + + raise ValueError(f"Transformer {self} cannot reverse {literal_type}") + def get_literal_type(self, t: Type[T]) -> LiteralType: return LiteralType( blob=_core_types.BlobType( diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 8463a4b6cd..6880fb4c67 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -362,6 +362,14 @@ def test_guessing_basic(): pt = TypeEngine.guess_python_type(lt) assert pt is type(None) + lt = model_types.LiteralType( + blob=BlobType( + format=FlytePickleTransformer.PYTHON_PICKLE_FORMAT, dimensionality=BlobType.BlobDimensionality.SINGLE + ) + ) + pt = TypeEngine.guess_python_type(lt) + assert pt is FlytePickle + def test_guessing_containers(): b = model_types.LiteralType(simple=model_types.SimpleType.BOOLEAN) @@ -778,6 +786,25 @@ def test_list_of_unions(): assert v == ["hello", 123, "world"] +def test_pickle_type(): + class Foo(object): + def __init__(self, number: int): + self.number = number + + lt = TypeEngine.to_literal_type(FlytePickle) + assert lt.blob.format == FlytePickleTransformer.PYTHON_PICKLE_FORMAT + assert lt.blob.dimensionality == BlobType.BlobDimensionality.SINGLE + + ctx = FlyteContextManager.current_context() + lv = TypeEngine.to_literal(ctx, Foo(1), FlytePickle, lt) + assert "/tmp/flyte/" in lv.scalar.blob.uri + + transformer = FlytePickleTransformer() + gt = transformer.guess_python_type(lt) + pv = transformer.to_python_value(ctx, lv, expected_python_type=gt) + assert Foo(1).number == pv.number + + def test_enum_in_dataclass(): @dataclass_json @dataclass From ac55099c752efbe1a5f106bde9fe904762f3f483 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Wed, 8 Dec 2021 02:10:00 +0800 Subject: [PATCH 030/128] Add support FlyteFile in dataclass (#725) * Add support Flyte File and directory in dataclass Signed-off-by: Kevin Su * Fixed tests Signed-off-by: Kevin Su * Fixed tests Signed-off-by: Kevin Su * Updated Signed-off-by: Kevin Su * Updated Signed-off-by: Kevin Su * Fixed tests Signed-off-by: Kevin Su --- flytekit/core/type_engine.py | 105 +++++++++++++++---- flytekit/types/directory/types.py | 26 +++-- flytekit/types/file/file.py | 34 +++--- tests/flytekit/unit/core/test_type_engine.py | 73 ++++++++++++- tests/flytekit/unit/core/test_type_hints.py | 71 ++++++++++++- 5 files changed, 257 insertions(+), 52 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 2550d6f73f..02dafc8f6b 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -29,7 +29,18 @@ from flytekit.models import interface as _interface_models from flytekit.models import types as _type_models from flytekit.models.core import types as _core_types -from flytekit.models.literals import Literal, LiteralCollection, LiteralMap, Primitive, Scalar, Schema, Union, Void +from flytekit.models.literals import ( + Blob, + BlobMetadata, + Literal, + LiteralCollection, + LiteralMap, + Primitive, + Scalar, + Schema, + Union, + Void, +) from flytekit.models.types import LiteralType, SimpleType, TypeStructure, UnionType try: @@ -289,29 +300,79 @@ def _serialize_flyte_type(self, python_val: T, python_type: Type[T]): """ If any field inside the dataclass is flyte type, we should use flyte type transformer for that field. """ - from flytekit.types.schema.types import FlyteSchema, FlyteSchemaTransformer + from flytekit.types.directory.types import FlyteDirectory + from flytekit.types.file import FlyteFile + from flytekit.types.schema.types import FlyteSchema for f in dataclasses.fields(python_type): v = python_val.__getattribute__(f.name) - if inspect.isclass(f.type) and issubclass(f.type, FlyteSchema): - FlyteSchemaTransformer().to_literal(FlyteContext.current_context(), v, f.type, None) + if inspect.isclass(f.type) and ( + issubclass(f.type, FlyteSchema) or issubclass(f.type, FlyteFile) or issubclass(f.type, FlyteDirectory) + ): + TypeEngine.to_literal(FlyteContext.current_context(), v, f.type, None) elif dataclasses.is_dataclass(f.type): self._serialize_flyte_type(v, f.type) - def _deserialize_flyte_type(self, python_val: T, expected_python_type: Type["FlyteSchema"]): + def _deserialize_flyte_type(self, python_val: T, expected_python_type: Type) -> T: + from flytekit.types.directory.types import FlyteDirectory, FlyteDirToMultipartBlobTransformer + from flytekit.types.file.file import FlyteFile, FlyteFilePathTransformer from flytekit.types.schema.types import FlyteSchema, FlyteSchemaTransformer - for f in dataclasses.fields(expected_python_type): - v = python_val.__getattribute__(f.name) - if inspect.isclass(f.type) and issubclass(f.type, FlyteSchema): - t = FlyteSchemaTransformer() - t.to_python_value( - FlyteContext.current_context(), - Literal(scalar=Scalar(schema=Schema(v.remote_path, t._get_schema_type(f.type)))), - f.type, - ) - elif dataclasses.is_dataclass(f.type): - self._deserialize_flyte_type(v, f.type) + if not dataclasses.is_dataclass(expected_python_type): + return python_val + + if issubclass(expected_python_type, FlyteSchema): + t = FlyteSchemaTransformer() + return t.to_python_value( + FlyteContext.current_context(), + Literal(scalar=Scalar(schema=Schema(python_val.remote_path, t._get_schema_type(expected_python_type)))), + expected_python_type, + ) + elif issubclass(expected_python_type, FlyteFile): + return FlyteFilePathTransformer().to_python_value( + FlyteContext.current_context(), + Literal( + scalar=Scalar( + blob=Blob( + metadata=BlobMetadata( + type=_core_types.BlobType( + format="", dimensionality=_core_types.BlobType.BlobDimensionality.SINGLE + ) + ), + uri=python_val.path, + ) + ) + ), + expected_python_type, + ) + elif issubclass(expected_python_type, FlyteDirectory): + return FlyteDirToMultipartBlobTransformer().to_python_value( + FlyteContext.current_context(), + Literal( + scalar=Scalar( + blob=Blob( + metadata=BlobMetadata( + type=_core_types.BlobType( + format="", dimensionality=_core_types.BlobType.BlobDimensionality.MULTIPART + ) + ), + uri=python_val.path, + ) + ) + ), + expected_python_type, + ) + else: + for f in dataclasses.fields(expected_python_type): + value = python_val.__getattribute__(f.name) + if hasattr(f.type, "__origin__") and f.type.__origin__ is list: + value = [self._deserialize_flyte_type(v, f.type.__args__[0]) for v in value] + elif hasattr(f.type, "__origin__") and f.type.__origin__ is dict: + value = {k: self._deserialize_flyte_type(v, f.type.__args__[1]) for k, v in value.items()} + else: + value = self._deserialize_flyte_type(value, f.type) + python_val.__setattr__(f.name, value) + return python_val def _fix_val_int(self, t: typing.Type, val: typing.Any) -> typing.Any: if t == int: @@ -357,8 +418,7 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: ) dc = cast(DataClassJsonMixin, expected_python_type).from_json(_json_format.MessageToJson(lv.scalar.generic)) - self._deserialize_flyte_type(dc, expected_python_type) - return self._fix_dataclass_int(expected_python_type, dc) + return self._fix_dataclass_int(expected_python_type, self._deserialize_flyte_type(dc, expected_python_type)) def guess_python_type(self, literal_type: LiteralType) -> Type[T]: if literal_type.simple == SimpleType.STRUCT: @@ -958,18 +1018,21 @@ def convert_json_schema_to_python_class(schema: dict, schema_name) -> Type[datac """ attribute_list = [] for property_key, property_val in schema[schema_name]["properties"].items(): + property_type = property_val["type"] # Handle list if property_val["type"] == "array": attribute_list.append((property_key, typing.List[_get_element_type(property_val["items"])])) # Handle dataclass and dict - elif property_val["type"] == "object": - if "$ref" in property_val: + elif property_type == "object": + if property_val.get("$ref"): name = property_val["$ref"].split("/")[-1] attribute_list.append((property_key, convert_json_schema_to_python_class(schema, name))) - else: + elif property_val.get("additionalProperties"): attribute_list.append( (property_key, typing.Dict[str, _get_element_type(property_val["additionalProperties"])]) ) + else: + attribute_list.append((property_key, typing.Dict[str, _get_element_type(property_val)])) # Handle int, float, bool or str else: attribute_list.append([property_key, _get_element_type(property_val)]) diff --git a/flytekit/types/directory/types.py b/flytekit/types/directory/types.py index 81f4fb0fd0..b08323cd1e 100644 --- a/flytekit/types/directory/types.py +++ b/flytekit/types/directory/types.py @@ -3,8 +3,12 @@ import os import pathlib import typing +from dataclasses import dataclass, field from pathlib import Path +from dataclasses_json import config, dataclass_json +from marshmallow import fields + from flytekit.core.context_manager import FlyteContext from flytekit.core.type_engine import TypeEngine, TypeTransformer from flytekit.models import types as _type_models @@ -19,7 +23,10 @@ def noop(): ... +@dataclass_json +@dataclass class FlyteDirectory(os.PathLike, typing.Generic[T]): + path: typing.Union[str, os.PathLike] = field(default=None, metadata=config(mm_field=fields.String())) """ .. warning:: @@ -106,15 +113,16 @@ def t1(in1: FlyteDirectory["svg"]): field in the ``BlobType``. """ - def __init__(self, path: str, downloader: typing.Callable = None, remote_directory=None): + def __init__(self, path: typing.Union[str, os.PathLike], downloader: typing.Callable = None, remote_directory=None): """ :param path: The source path that users are expected to call open() on :param downloader: Optional function that can be passed that used to delay downloading of the actual fil until a user actually calls open(). :param remote_directory: If the user wants to return something and also specify where it should be uploaded to. """ - - self._path = path + # Make this field public, so that the dataclass transformer can set a value for it + # https://github.com/flyteorg/flytekit/blob/bcc8541bd6227b532f8462563fe8aac902242b21/flytekit/core/type_engine.py#L298 + self.path = path self._downloader = downloader or noop self._downloaded = False self._remote_directory = remote_directory @@ -127,7 +135,7 @@ def __fspath__(self): if not self._downloaded: self._downloader() self._downloaded = True - return self._path + return self.path @classmethod def extension(cls) -> str: @@ -159,10 +167,6 @@ def downloaded(self) -> bool: def remote_directory(self) -> typing.Optional[str]: return self._remote_directory - @property - def path(self) -> str: - return self._path - @property def remote_source(self) -> str: """ @@ -175,10 +179,10 @@ def download(self) -> str: return self.__fspath__() def __repr__(self): - return self._path + return self.path def __str__(self): - return self._path + return self.path class FlyteDirToMultipartBlobTransformer(TypeTransformer[FlyteDirectory]): @@ -238,7 +242,7 @@ def to_literal( return Literal(scalar=Scalar(blob=Blob(metadata=meta, uri=python_val._remote_source))) source_path = python_val.path - # If the user specified the remote_path to be False, that means no matter what, do not upload. Also if the + # If the user specified the remote_directory to be False, that means no matter what, do not upload. Also if the # path given is already a remote path, say https://www.google.com, the concept of uploading to the Flyte # blob store doesn't make sense. if python_val.remote_directory is False or ctx.file_access.is_remote(source_path): diff --git a/flytekit/types/file/file.py b/flytekit/types/file/file.py index 826bc3916d..b08bcd565c 100644 --- a/flytekit/types/file/file.py +++ b/flytekit/types/file/file.py @@ -3,6 +3,10 @@ import os import pathlib import typing +from dataclasses import dataclass, field + +from dataclasses_json import config, dataclass_json +from marshmallow import fields from flytekit.core.context_manager import FlyteContext from flytekit.core.type_engine import TypeEngine, TypeTransformer, TypeTransformerFailedError @@ -20,7 +24,10 @@ def noop(): T = typing.TypeVar("T") +@dataclass_json +@dataclass class FlyteFile(os.PathLike, typing.Generic[T]): + path: typing.Union[str, os.PathLike] = field(default=None, metadata=config(mm_field=fields.String())) """ Since there is no native Python implementation of files and directories for the Flyte Blob type, (like how int exists for Flyte's Integer type) we need to create one so that users can express that their tasks take @@ -156,14 +163,18 @@ def extension(cls) -> str: return _SpecificFormatClass - def __init__(self, path: str, downloader: typing.Callable = noop, remote_path=None): + def __init__( + self, path: typing.Union[str, os.PathLike], downloader: typing.Callable = noop, remote_path: os.PathLike = None + ): """ :param path: The source path that users are expected to call open() on :param downloader: Optional function that can be passed that used to delay downloading of the actual fil until a user actually calls open(). :param remote_path: If the user wants to return something and also specify where it should be uploaded to. """ - self._path = path + # Make this field public, so that the dataclass transformer can set a value for it + # https://github.com/flyteorg/flytekit/blob/bcc8541bd6227b532f8462563fe8aac902242b21/flytekit/core/type_engine.py#L298 + self.path = path self._downloader = downloader self._downloaded = False self._remote_path = remote_path @@ -174,30 +185,26 @@ def __fspath__(self): if not self._downloaded: self._downloader() self._downloaded = True - return self._path + return self.path def __eq__(self, other): if isinstance(other, FlyteFile): return ( - self._path == other._path + self.path == other.path and self._remote_path == other._remote_path and self.extension() == other.extension() ) else: - return self._path == other + return self.path == other @property def downloaded(self) -> bool: return self._downloaded @property - def remote_path(self) -> typing.Optional[str]: + def remote_path(self) -> os.PathLike: return self._remote_path - @property - def path(self) -> str: - return self._path - @property def remote_source(self) -> str: """ @@ -210,10 +217,10 @@ def download(self) -> str: return self.__fspath__() def __repr__(self): - return self._path + return self.path def __str__(self): - return self._path + return self.path class FlyteFilePathTransformer(TypeTransformer[FlyteFile]): @@ -317,11 +324,10 @@ def to_literal( return Literal(scalar=Scalar(blob=Blob(metadata=meta, uri=source_path))) def to_python_value( - self, ctx: FlyteContext, lv: Literal, expected_python_type: typing.Union[typing.Type[FlyteFile]] + self, ctx: FlyteContext, lv: Literal, expected_python_type: typing.Union[typing.Type[FlyteFile], os.PathLike] ) -> FlyteFile: uri = lv.scalar.blob.uri - # In this condition, we still return a FlyteFile instance, but it's a simple one that has no downloading tricks # Using is instead of issubclass because FlyteFile does actually subclass it if expected_python_type is os.PathLike: diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 6880fb4c67..0768b07db6 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -1,5 +1,6 @@ import datetime import os +import tempfile import typing from dataclasses import asdict, dataclass from datetime import timedelta @@ -35,9 +36,10 @@ from flytekit.models.core.types import BlobType from flytekit.models.literals import Blob, BlobMetadata, Literal, LiteralCollection, LiteralMap, Primitive, Scalar, Void from flytekit.models.types import LiteralType, SimpleType, TypeStructure +from flytekit.types.directory import TensorboardLogs from flytekit.types.directory.types import FlyteDirectory from flytekit.types.file import JPEGImageFile -from flytekit.types.file.file import FlyteFile, FlyteFilePathTransformer +from flytekit.types.file.file import FlyteFile, FlyteFilePathTransformer, noop from flytekit.types.pickle import FlytePickle from flytekit.types.pickle.pickle import FlytePickleTransformer from flytekit.types.schema import FlyteSchema @@ -471,7 +473,6 @@ def test_dataclass_transformer(): }, }, } - tf = DataclassTransformer() t = tf.get_literal_type(TestStruct) assert t is not None @@ -521,6 +522,74 @@ def test_dataclass_int_preserving(): assert ot == o +def test_flyte_file_in_dataclass(): + @dataclass_json + @dataclass + class TestInnerFileStruct(object): + a: JPEGImageFile + b: typing.List[FlyteFile] + c: typing.Dict[str, FlyteFile] + + @dataclass_json + @dataclass + class TestFileStruct(object): + a: FlyteFile + b: TestInnerFileStruct + + f = FlyteFile("s3://tmp/file") + o = TestFileStruct(a=f, b=TestInnerFileStruct(a=JPEGImageFile("s3://tmp/file.jpeg"), b=[f], c={"hello": f})) + + ctx = FlyteContext.current_context() + tf = DataclassTransformer() + lt = tf.get_literal_type(TestFileStruct) + lv = tf.to_literal(ctx, o, TestFileStruct, lt) + ot = tf.to_python_value(ctx, lv=lv, expected_python_type=TestFileStruct) + assert ot.a._downloader is not noop + assert ot.b.a._downloader is not noop + assert ot.b.b[0]._downloader is not noop + assert ot.b.c["hello"]._downloader is not noop + + assert o.a.path == ot.a.remote_source + assert o.b.a.path == ot.b.a.remote_source + assert o.b.b[0].path == ot.b.b[0].remote_source + assert o.b.c["hello"].path == ot.b.c["hello"].remote_source + + +def test_flyte_directory_in_dataclass(): + @dataclass_json + @dataclass + class TestInnerFileStruct(object): + a: TensorboardLogs + b: typing.List[FlyteDirectory] + c: typing.Dict[str, FlyteDirectory] + + @dataclass_json + @dataclass + class TestFileStruct(object): + a: FlyteDirectory + b: TestInnerFileStruct + + tempdir = tempfile.mkdtemp(prefix="flyte-") + f = FlyteDirectory(tempdir) + o = TestFileStruct(a=f, b=TestInnerFileStruct(a=TensorboardLogs("s3://tensorboard"), b=[f], c={"hello": f})) + + ctx = FlyteContext.current_context() + tf = DataclassTransformer() + lt = tf.get_literal_type(TestFileStruct) + lv = tf.to_literal(ctx, o, TestFileStruct, lt) + ot = tf.to_python_value(ctx, lv=lv, expected_python_type=TestFileStruct) + + assert ot.a._downloader is not noop + assert ot.b.a._downloader is not noop + assert ot.b.b[0]._downloader is not noop + assert ot.b.c["hello"]._downloader is not noop + + assert o.a.path == ot.a.path + assert o.b.a.path == ot.b.a.remote_source + assert o.b.b[0].path == ot.b.b[0].path + assert o.b.c["hello"].path == ot.b.c["hello"].path + + # Enums should have string values class Color(Enum): RED = "red" diff --git a/tests/flytekit/unit/core/test_type_hints.py b/tests/flytekit/unit/core/test_type_hints.py index 57de0c29b2..131bd7cc09 100644 --- a/tests/flytekit/unit/core/test_type_hints.py +++ b/tests/flytekit/unit/core/test_type_hints.py @@ -35,7 +35,8 @@ from flytekit.models.interface import Parameter from flytekit.models.task import Resources as _resource_models from flytekit.models.types import LiteralType, SimpleType -from flytekit.types.file import FlyteFile +from flytekit.types.directory import FlyteDirectory, TensorboardLogs +from flytekit.types.file import FlyteFile, PNGImageFile from flytekit.types.schema import FlyteSchema, SchemaOpenMode serialization_settings = context_manager.SerializationSettings( @@ -352,6 +353,68 @@ def test_user_demo_test(mock_sql): assert context_manager.FlyteContextManager.size() == 1 +def test_flyte_file_in_dataclass(): + @dataclass_json + @dataclass + class InnerFileStruct(object): + a: FlyteFile + b: PNGImageFile + + @dataclass_json + @dataclass + class FileStruct(object): + a: FlyteFile + b: InnerFileStruct + + @task + def t1(path: str) -> FileStruct: + file = FlyteFile(path) + fs = FileStruct(a=file, b=InnerFileStruct(a=file, b=PNGImageFile(path))) + return fs + + @task + def t2(fs: FileStruct) -> os.PathLike: + return fs.a.path + + @workflow + def wf(path: str) -> os.PathLike: + n1 = t1(path=path) + return t2(fs=n1) + + assert "/tmp/flyte/" in wf(path="s3://somewhere").path + + +def test_flyte_directory_in_dataclass(): + @dataclass_json + @dataclass + class InnerFileStruct(object): + a: FlyteDirectory + b: TensorboardLogs + + @dataclass_json + @dataclass + class FileStruct(object): + a: FlyteDirectory + b: InnerFileStruct + + @task + def t1(path: str) -> FileStruct: + dir = FlyteDirectory(path) + fs = FileStruct(a=dir, b=InnerFileStruct(a=dir, b=TensorboardLogs(path))) + return fs + + @task + def t2(fs: FileStruct) -> os.PathLike: + return fs.a.path + + @workflow + def wf(path: str) -> os.PathLike: + n1 = t1(path=path) + return t2(fs=n1) + + assert "/tmp/flyte/" in wf(path="s3://somewhere").path + + def test_wf1_with_map(): @task def t1(a: int) -> int: @@ -1433,7 +1496,7 @@ class Foo(object): @dataclass class Bar(object): x: int - y: str + y: dict z: Foo @task @@ -1452,14 +1515,14 @@ def t1() -> Foo: @task def t2() -> Bar: - return Bar(x=1, y="bar", z=Foo(x=1, y="foo", z={"hello": "world"})) + return Bar(x=1, y={"hello": "world"}, z=Foo(x=1, y="foo", z={"hello": "world"})) task_spec = get_serializable(OrderedDict(), serialization_settings, t2) pt_map = TypeEngine.guess_python_types(task_spec.template.interface.outputs) assert dataclasses.is_dataclass(pt_map["o0"]) output_lm = t2.dispatch_execute(ctx, _literal_models.LiteralMap(literals={})) - expected_struct.update({"x": 1, "y": "bar", "z": {"x": 1, "y": "foo", "z": {"hello": "world"}}}) + expected_struct.update({"x": 1, "y": {"hello": "world"}, "z": {"x": 1, "y": "foo", "z": {"hello": "world"}}}) assert output_lm.literals["o0"].scalar.generic == expected_struct From 4411b1cad75da32155a7e01d0f6a373f07cca9a6 Mon Sep 17 00:00:00 2001 From: Niels Bantilan Date: Tue, 7 Dec 2021 19:26:51 -0500 Subject: [PATCH 031/128] add task_resolver arg to @task decorator (#765) Signed-off-by: Niels Bantilan Signed-off-by: maximsmol --- flytekit/core/task.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/flytekit/core/task.py b/flytekit/core/task.py index 24fe283399..62f982a11b 100644 --- a/flytekit/core/task.py +++ b/flytekit/core/task.py @@ -2,7 +2,7 @@ import inspect from typing import Any, Callable, Dict, List, Optional, Type, Union -from flytekit.core.base_task import TaskMetadata +from flytekit.core.base_task import TaskMetadata, TaskResolverMixin from flytekit.core.interface import transform_signature_to_interface from flytekit.core.python_function_task import PythonFunctionTask from flytekit.core.reference_entity import ReferenceEntity, TaskReference @@ -87,6 +87,7 @@ def task( limits: Optional[Resources] = None, secret_requests: Optional[List[Secret]] = None, execution_mode: Optional[PythonFunctionTask.ExecutionBehavior] = PythonFunctionTask.ExecutionBehavior.DEFAULT, + task_resolver: Optional[TaskResolverMixin] = None, ) -> Union[Callable, PythonFunctionTask]: """ This is the core decorator to use for any task type in flytekit. @@ -170,6 +171,7 @@ def foo2(): Refer to :py:class:`Secret` to understand how to specify the request for a secret. It may change based on the backend provider. :param execution_mode: This is mainly for internal use. Please ignore. It is filled in automatically. + :param task_resolver: Provide a custom task resolver. """ def wrapper(fn) -> PythonFunctionTask: @@ -192,6 +194,7 @@ def wrapper(fn) -> PythonFunctionTask: limits=limits, secret_requests=secret_requests, execution_mode=execution_mode, + task_resolver=task_resolver, ) return task_instance From b89124f3d0a18d67d6c255dbe3a084b74935e1f6 Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Tue, 7 Dec 2021 16:51:30 -0800 Subject: [PATCH 032/128] Copy metadata into map task from underlying (#766) Signed-off-by: Yee Hing Tong Signed-off-by: maximsmol --- flytekit/core/map_task.py | 2 ++ tests/flytekit/unit/core/test_map_task.py | 14 ++++++++++++++ 2 files changed, 16 insertions(+) diff --git a/flytekit/core/map_task.py b/flytekit/core/map_task.py index fa48b474ca..60061cbb0e 100644 --- a/flytekit/core/map_task.py +++ b/flytekit/core/map_task.py @@ -58,6 +58,8 @@ def __init__( self._max_concurrency = concurrency self._min_success_ratio = min_success_ratio self._array_task_interface = python_function_task.python_interface + if "metadata" not in kwargs and python_function_task.metadata: + kwargs["metadata"] = python_function_task.metadata super().__init__( name=name, interface=collection_interface, diff --git a/tests/flytekit/unit/core/test_map_task.py b/tests/flytekit/unit/core/test_map_task.py index 80ab499416..95df669829 100644 --- a/tests/flytekit/unit/core/test_map_task.py +++ b/tests/flytekit/unit/core/test_map_task.py @@ -18,6 +18,12 @@ def t1(a: int) -> str: return str(b) +@task(cache=True, cache_version="1") +def t2(a: int) -> str: + b = a + 2 + return str(b) + + # This test is for documentation. def test_map_docs(): # test_map_task_start @@ -162,3 +168,11 @@ def many_outputs(a: int) -> (int, str): with pytest.raises(ValueError): _ = map_task(many_inputs) + + +def test_map_task_metadata(): + map_meta = TaskMetadata(retries=1) + mapped_1 = map_task(t2, metadata=map_meta) + assert mapped_1.metadata is map_meta + mapped_2 = map_task(t2) + assert mapped_2.metadata is t2.metadata From ec13d835c45df9944b88fa8e51d9ff24be4d4738 Mon Sep 17 00:00:00 2001 From: Stef Nelson-Lindall Date: Thu, 9 Dec 2021 10:41:13 -0800 Subject: [PATCH 033/128] Support for delayed annotations (#760) Signed-off-by: Stefan Nelson-Lindall Signed-off-by: Yee Hing Tong Signed-off-by: maximsmol --- flytekit/core/interface.py | 16 ++++-- flytekit/core/launch_plan.py | 5 +- flytekit/core/promise.py | 2 +- flytekit/core/python_function_task.py | 7 +-- flytekit/core/task.py | 5 +- flytekit/core/type_engine.py | 7 ++- flytekit/core/workflow.py | 7 +-- .../unit/core/functools/simple_decorator.py | 1 + .../unit/core/functools/test_decorators.py | 1 + tests/flytekit/unit/core/test_interface.py | 57 ++++++++++--------- tests/flytekit/unit/core/test_type_delayed.py | 27 +++++++++ 11 files changed, 86 insertions(+), 49 deletions(-) create mode 100644 tests/flytekit/unit/core/test_type_delayed.py diff --git a/flytekit/core/interface.py b/flytekit/core/interface.py index 4ef55b809c..216e337282 100644 --- a/flytekit/core/interface.py +++ b/flytekit/core/interface.py @@ -267,19 +267,24 @@ def _change_unrecognized_type_to_pickle(t: Type[T]) -> Type[T]: return t -def transform_signature_to_interface(signature: inspect.Signature, docstring: Optional[Docstring] = None) -> Interface: +def transform_function_to_interface(fn: Callable, docstring: Optional[Docstring] = None) -> Interface: """ From the annotations on a task function that the user should have provided, and the output names they want to use for each output parameter, construct the TypedInterface object For now the fancy object, maybe in the future a dumb object. + """ - outputs = extract_return_annotation(signature.return_annotation) + type_hints = typing.get_type_hints(fn) + signature = inspect.signature(fn) + return_annotation = type_hints.get("return", None) + + outputs = extract_return_annotation(return_annotation) for k, v in outputs.items(): outputs[k] = _change_unrecognized_type_to_pickle(v) inputs = OrderedDict() for k, v in signature.parameters.items(): - annotation = v.annotation + annotation = type_hints.get(k, None) 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] = (_change_unrecognized_type_to_pickle(annotation), default) @@ -287,7 +292,6 @@ def transform_signature_to_interface(signature: inspect.Signature, docstring: Op # This is just for typing.NamedTuples - in those cases, the user can select a name to call the NamedTuple. We # would like to preserve that name in our custom collections.namedtuple. custom_name = None - return_annotation = signature.return_annotation if hasattr(return_annotation, "__bases__"): bases = return_annotation.__bases__ if len(bases) == 1 and bases[0] == tuple and hasattr(return_annotation, "_fields"): @@ -334,7 +338,7 @@ def output_name_generator(length: int) -> Generator[str, None, None]: yield default_output_name(x) -def extract_return_annotation(return_annotation: Union[Type, Tuple]) -> Dict[str, Type]: +def extract_return_annotation(return_annotation: Union[Type, Tuple, None]) -> Dict[str, Type]: """ The purpose of this function is to sort out whether a function is returning one thing, or multiple things, and to name the outputs accordingly, either by using our default name function, or from a typing.NamedTuple. @@ -368,7 +372,7 @@ def t(a: int, b: str) -> Dict[str, int]: ... # Handle Option 6 # We can think about whether we should add a default output name with type None in the future. - if return_annotation is None or return_annotation is inspect.Signature.empty: + if return_annotation in (None, type(None), inspect.Signature.empty): return {} # This statement results in true for typing.Namedtuple, single and void return types, so this diff --git a/flytekit/core/launch_plan.py b/flytekit/core/launch_plan.py index 481871977e..217db39652 100644 --- a/flytekit/core/launch_plan.py +++ b/flytekit/core/launch_plan.py @@ -1,11 +1,10 @@ from __future__ import annotations -import inspect from typing import Any, Callable, Dict, List, Optional, Type from flytekit.core import workflow as _annotated_workflow from flytekit.core.context_manager import FlyteContext, FlyteContextManager, FlyteEntities -from flytekit.core.interface import Interface, transform_inputs_to_parameters, transform_signature_to_interface +from flytekit.core.interface import Interface, transform_function_to_interface, transform_inputs_to_parameters from flytekit.core.promise import create_and_link_node, translate_inputs_to_literals from flytekit.core.reference_entity import LaunchPlanReference, ReferenceEntity from flytekit.models import common as _common_models @@ -399,7 +398,7 @@ def reference_launch_plan( """ def wrapper(fn) -> ReferenceLaunchPlan: - interface = transform_signature_to_interface(inspect.signature(fn)) + interface = transform_function_to_interface(fn) return ReferenceLaunchPlan(project, domain, name, version, interface.inputs, interface.outputs) return wrapper diff --git a/flytekit/core/promise.py b/flytekit/core/promise.py index 01a19cca1e..7c9b780e19 100644 --- a/flytekit/core/promise.py +++ b/flytekit/core/promise.py @@ -500,7 +500,7 @@ def create_task_output( if len(promises) == 1: if not entity_interface: return promises[0] - # See transform_signature_to_interface for more information, we're using the existence of a name as a proxy + # See transform_function_to_interface for more information, we're using the existence of a name as a proxy # for the user having specified a one-element typing.NamedTuple, which means we should _not_ extract it. We # should still return a tuple but it should be one of ours. if not entity_interface.output_tuple_name: diff --git a/flytekit/core/python_function_task.py b/flytekit/core/python_function_task.py index addea248f4..25a363b070 100644 --- a/flytekit/core/python_function_task.py +++ b/flytekit/core/python_function_task.py @@ -14,7 +14,6 @@ """ -import inspect from abc import ABC from collections import OrderedDict from enum import Enum @@ -24,7 +23,7 @@ from flytekit.core.base_task import Task, TaskResolverMixin from flytekit.core.context_manager import ExecutionState, FastSerializationSettings, FlyteContext, FlyteContextManager from flytekit.core.docstring import Docstring -from flytekit.core.interface import transform_signature_to_interface +from flytekit.core.interface import transform_function_to_interface from flytekit.core.python_auto_container import PythonAutoContainerTask, default_task_resolver from flytekit.core.tracker import is_functools_wrapped_module_level, isnested, istestfunction from flytekit.core.workflow import ( @@ -114,9 +113,7 @@ def __init__( """ if task_function is None: raise ValueError("TaskFunction is a required parameter for PythonFunctionTask") - self._native_interface = transform_signature_to_interface( - inspect.signature(task_function), Docstring(callable_=task_function) - ) + self._native_interface = transform_function_to_interface(task_function, Docstring(callable_=task_function)) mutated_interface = self._native_interface.remove_inputs(ignore_input_vars) super().__init__( task_type=task_type, diff --git a/flytekit/core/task.py b/flytekit/core/task.py index 62f982a11b..45b7e2dc94 100644 --- a/flytekit/core/task.py +++ b/flytekit/core/task.py @@ -1,9 +1,8 @@ import datetime as _datetime -import inspect from typing import Any, Callable, Dict, List, Optional, Type, Union from flytekit.core.base_task import TaskMetadata, TaskResolverMixin -from flytekit.core.interface import transform_signature_to_interface +from flytekit.core.interface import transform_function_to_interface from flytekit.core.python_function_task import PythonFunctionTask from flytekit.core.reference_entity import ReferenceEntity, TaskReference from flytekit.core.resources import Resources @@ -240,7 +239,7 @@ def reference_task( """ def wrapper(fn) -> ReferenceTask: - interface = transform_signature_to_interface(inspect.signature(fn)) + interface = transform_function_to_interface(fn) return ReferenceTask(project, domain, name, version, interface.inputs, interface.outputs) return wrapper diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 02dafc8f6b..178b550136 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -277,7 +277,12 @@ def get_literal_type(self, t: Type[T]) -> LiteralType: v.load_by = LoadDumpOptions.name schema = JSONSchema().dump(s) except Exception as e: - logger.warn("failed to extract schema for object %s, (will run schemaless) error: %s", str(t), e) + # https://github.com/lovasoa/marshmallow_dataclass/issues/13 + logger.warning( + f"Failed to extract schema for object {t}, (will run schemaless) error: {e}" + f"If you have postponed annotations turned on (PEP 563) turn it off please. Postponed" + f"evaluation doesn't work with json dataclasses" + ) return _primitives.Generic.to_flyte_literal_type(metadata=schema) diff --git a/flytekit/core/workflow.py b/flytekit/core/workflow.py index a14be3d0a5..744ecfbb11 100644 --- a/flytekit/core/workflow.py +++ b/flytekit/core/workflow.py @@ -1,6 +1,5 @@ from __future__ import annotations -import inspect from dataclasses import dataclass from enum import Enum from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union @@ -15,9 +14,9 @@ from flytekit.core.docstring import Docstring from flytekit.core.interface import ( Interface, + transform_function_to_interface, transform_inputs_to_parameters, transform_interface_to_typed_interface, - transform_signature_to_interface, ) from flytekit.core.launch_plan import LaunchPlan from flytekit.core.node import Node @@ -574,7 +573,7 @@ def __init__( ): name = f"{workflow_function.__module__}.{workflow_function.__name__}" self._workflow_function = workflow_function - native_interface = transform_signature_to_interface(inspect.signature(workflow_function), docstring=docstring) + native_interface = transform_function_to_interface(workflow_function, docstring=docstring) # TODO do we need this - can this not be in launchplan only? # This can be in launch plan only, but is here only so that we don't have to re-evaluate. Or @@ -770,7 +769,7 @@ def reference_workflow( """ def wrapper(fn) -> ReferenceWorkflow: - interface = transform_signature_to_interface(inspect.signature(fn)) + interface = transform_function_to_interface(fn) return ReferenceWorkflow(project, domain, name, version, interface.inputs, interface.outputs) return wrapper diff --git a/tests/flytekit/unit/core/functools/simple_decorator.py b/tests/flytekit/unit/core/functools/simple_decorator.py index cf568c5aa9..a51a283be5 100644 --- a/tests/flytekit/unit/core/functools/simple_decorator.py +++ b/tests/flytekit/unit/core/functools/simple_decorator.py @@ -1,4 +1,5 @@ """Script used for testing local execution of functool.wraps-wrapped tasks""" +from __future__ import annotations import os from functools import wraps diff --git a/tests/flytekit/unit/core/functools/test_decorators.py b/tests/flytekit/unit/core/functools/test_decorators.py index f87e714dc9..4e0bb5818b 100644 --- a/tests/flytekit/unit/core/functools/test_decorators.py +++ b/tests/flytekit/unit/core/functools/test_decorators.py @@ -1,4 +1,5 @@ """Test local execution of files that use functools to decorate tasks and workflows.""" +from __future__ import annotations import os import subprocess diff --git a/tests/flytekit/unit/core/test_interface.py b/tests/flytekit/unit/core/test_interface.py index 81a01517de..8e55ee1bb4 100644 --- a/tests/flytekit/unit/core/test_interface.py +++ b/tests/flytekit/unit/core/test_interface.py @@ -1,4 +1,3 @@ -import inspect import os import typing from typing import Dict, List @@ -7,9 +6,9 @@ from flytekit.core.docstring import Docstring from flytekit.core.interface import ( extract_return_annotation, + transform_function_to_interface, transform_inputs_to_parameters, transform_interface_to_typed_interface, - transform_signature_to_interface, transform_variable_map, ) from flytekit.models.core import types as _core_types @@ -21,7 +20,7 @@ def test_extract_only(): def x() -> typing.NamedTuple("NT1", x_str=str, y_int=int): ... - return_types = extract_return_annotation(inspect.signature(x).return_annotation) + return_types = extract_return_annotation(typing.get_type_hints(x).get("return", None)) assert len(return_types) == 2 assert return_types["x_str"] == str assert return_types["y_int"] == int @@ -29,7 +28,7 @@ def x() -> typing.NamedTuple("NT1", x_str=str, y_int=int): def t() -> List[int]: ... - return_type = extract_return_annotation(inspect.signature(t).return_annotation) + return_type = extract_return_annotation(typing.get_type_hints(t).get("return", None)) assert len(return_type) == 1 assert return_type["o0"]._name == "List" assert return_type["o0"].__origin__ == list @@ -37,7 +36,7 @@ def t() -> List[int]: def t() -> Dict[str, int]: ... - return_type = extract_return_annotation(inspect.signature(t).return_annotation) + return_type = extract_return_annotation(typing.get_type_hints(t).get("return", None)) assert len(return_type) == 1 assert return_type["o0"]._name == "Dict" assert return_type["o0"].__origin__ == dict @@ -45,7 +44,7 @@ def t() -> Dict[str, int]: def t(a: int, b: str) -> typing.Tuple[int, str]: ... - return_type = extract_return_annotation(inspect.signature(t).return_annotation) + return_type = extract_return_annotation(typing.get_type_hints(t).get("return", None)) assert len(return_type) == 2 assert return_type["o0"] == int assert return_type["o1"] == str @@ -53,7 +52,7 @@ def t(a: int, b: str) -> typing.Tuple[int, str]: def t(a: int, b: str) -> (int, str): ... - return_type = extract_return_annotation(inspect.signature(t).return_annotation) + return_type = extract_return_annotation(typing.get_type_hints(t).get("return", None)) assert len(return_type) == 2 assert return_type["o0"] == int assert return_type["o1"] == str @@ -61,27 +60,33 @@ def t(a: int, b: str) -> (int, str): def t(a: int, b: str) -> str: ... - return_type = extract_return_annotation(inspect.signature(t).return_annotation) + return_type = extract_return_annotation(typing.get_type_hints(t).get("return", None)) assert len(return_type) == 1 assert return_type["o0"] == str + def t(a: int, b: str): + ... + + return_type = extract_return_annotation(typing.get_type_hints(t).get("return", None)) + assert len(return_type) == 0 + def t(a: int, b: str) -> None: ... - return_type = extract_return_annotation(inspect.signature(t).return_annotation) + return_type = extract_return_annotation(typing.get_type_hints(t).get("return", None)) assert len(return_type) == 0 def t(a: int, b: str) -> List[int]: ... - return_type = extract_return_annotation(inspect.signature(t).return_annotation) + return_type = extract_return_annotation(typing.get_type_hints(t).get("return", None)) assert len(return_type) == 1 assert return_type["o0"] == List[int] def t(a: int, b: str) -> Dict[str, int]: ... - return_type = extract_return_annotation(inspect.signature(t).return_annotation) + return_type = extract_return_annotation(typing.get_type_hints(t).get("return", None)) assert len(return_type) == 1 assert return_type["o0"] == Dict[str, int] @@ -95,11 +100,11 @@ def x(a: int, b: str) -> typing.NamedTuple("NT1", x_str=str, y_int=int): def y(a: int, b: str) -> nt1: return nt1("hello world", 5) - result = transform_variable_map(extract_return_annotation(inspect.signature(x).return_annotation)) + result = transform_variable_map(extract_return_annotation(typing.get_type_hints(x).get("return", None))) assert result["x_str"].type.simple == 3 assert result["y_int"].type.simple == 1 - result = transform_variable_map(extract_return_annotation(inspect.signature(y).return_annotation)) + result = transform_variable_map(extract_return_annotation(typing.get_type_hints(y).get("return", None))) assert result["x_str"].type.simple == 3 assert result["y_int"].type.simple == 1 @@ -108,7 +113,7 @@ def test_unnamed_typing_tuple(): def z(a: int, b: str) -> typing.Tuple[int, str]: return 5, "hello world" - result = transform_variable_map(extract_return_annotation(inspect.signature(z).return_annotation)) + result = transform_variable_map(extract_return_annotation(typing.get_type_hints(z).get("return", None))) assert result["o0"].type.simple == 1 assert result["o1"].type.simple == 3 @@ -117,7 +122,7 @@ def test_regular_tuple(): def q(a: int, b: str) -> (int, str): return 5, "hello world" - result = transform_variable_map(extract_return_annotation(inspect.signature(q).return_annotation)) + result = transform_variable_map(extract_return_annotation(typing.get_type_hints(q).get("return", None))) assert result["o0"].type.simple == 1 assert result["o1"].type.simple == 3 @@ -126,7 +131,7 @@ def test_single_output_new_decorator(): def q(a: int, b: str) -> int: return a + len(b) - result = transform_variable_map(extract_return_annotation(inspect.signature(q).return_annotation)) + result = transform_variable_map(extract_return_annotation(typing.get_type_hints(q).get("return", None))) assert result["o0"].type.simple == 1 @@ -134,7 +139,7 @@ def test_sig_files(): def q() -> os.PathLike: ... - result = transform_variable_map(extract_return_annotation(inspect.signature(q).return_annotation)) + result = transform_variable_map(extract_return_annotation(typing.get_type_hints(q).get("return", None))) assert isinstance(result["o0"].type.blob, _core_types.BlobType) @@ -142,7 +147,7 @@ def test_file_types(): def t1() -> FlyteFile[typing.TypeVar("svg")]: ... - return_type = extract_return_annotation(inspect.signature(t1).return_annotation) + return_type = extract_return_annotation(typing.get_type_hints(t1).get("return", None)) assert return_type["o0"].extension() == FlyteFile[typing.TypeVar("svg")].extension() @@ -152,7 +157,7 @@ def test_parameters_and_defaults(): def z(a: int, b: str) -> typing.Tuple[int, str]: ... - our_interface = transform_signature_to_interface(inspect.signature(z)) + our_interface = transform_function_to_interface(z) params = transform_inputs_to_parameters(ctx, our_interface) assert params.parameters["a"].required assert params.parameters["a"].default is None @@ -162,7 +167,7 @@ def z(a: int, b: str) -> typing.Tuple[int, str]: def z(a: int, b: str = "hello") -> typing.Tuple[int, str]: ... - our_interface = transform_signature_to_interface(inspect.signature(z)) + our_interface = transform_function_to_interface(z) params = transform_inputs_to_parameters(ctx, our_interface) assert params.parameters["a"].required assert params.parameters["a"].default is None @@ -172,7 +177,7 @@ def z(a: int, b: str = "hello") -> typing.Tuple[int, str]: def z(a: int = 7, b: str = "eleven") -> typing.Tuple[int, str]: ... - our_interface = transform_signature_to_interface(inspect.signature(z)) + our_interface = transform_function_to_interface(z) params = transform_inputs_to_parameters(ctx, our_interface) assert not params.parameters["a"].required assert params.parameters["a"].default.scalar.primitive.integer == 7 @@ -193,7 +198,7 @@ def z(a: int, b: str) -> typing.Tuple[int, str]: """ ... - our_interface = transform_signature_to_interface(inspect.signature(z), Docstring(callable_=z)) + our_interface = transform_function_to_interface(z, Docstring(callable_=z)) params = transform_inputs_to_parameters(ctx, our_interface) assert params.parameters["a"].var.description == "foo" assert params.parameters["b"].var.description == "bar" @@ -211,7 +216,7 @@ def z(a: int, b: str) -> typing.Tuple[int, str]: """ ... - our_interface = transform_signature_to_interface(inspect.signature(z), Docstring(callable_=z)) + our_interface = transform_function_to_interface(z, Docstring(callable_=z)) typed_interface = transform_interface_to_typed_interface(our_interface) assert typed_interface.inputs.get("a").description == "foo" assert typed_interface.inputs.get("b").description == "bar" @@ -236,7 +241,7 @@ def z(a: int, b: str) -> typing.Tuple[int, str]: """ ... - our_interface = transform_signature_to_interface(inspect.signature(z), Docstring(callable_=z)) + our_interface = transform_function_to_interface(z, Docstring(callable_=z)) typed_interface = transform_interface_to_typed_interface(our_interface) assert typed_interface.inputs.get("a").description == "foo" assert typed_interface.inputs.get("b").description == "bar" @@ -264,7 +269,7 @@ def z(a: int, b: str) -> typing.NamedTuple("NT", x_str=str, y_int=int): """ ... - our_interface = transform_signature_to_interface(inspect.signature(z), Docstring(callable_=z)) + our_interface = transform_function_to_interface(z, Docstring(callable_=z)) typed_interface = transform_interface_to_typed_interface(our_interface) assert typed_interface.inputs.get("a").description == "foo" assert typed_interface.inputs.get("b").description == "bar" @@ -282,7 +287,7 @@ def __init__(self, name): def z(a: Foo) -> Foo: ... - our_interface = transform_signature_to_interface(inspect.signature(z)) + our_interface = transform_function_to_interface(z) params = transform_inputs_to_parameters(ctx, our_interface) assert params.parameters["a"].required assert params.parameters["a"].default is None diff --git a/tests/flytekit/unit/core/test_type_delayed.py b/tests/flytekit/unit/core/test_type_delayed.py new file mode 100644 index 0000000000..87daa91f47 --- /dev/null +++ b/tests/flytekit/unit/core/test_type_delayed.py @@ -0,0 +1,27 @@ +from __future__ import annotations + +import typing +from dataclasses import dataclass + +from dataclasses_json import dataclass_json + +from flytekit.core.type_engine import TypeEngine + + +@dataclass_json +@dataclass +class Foo(object): + x: int + y: str + z: typing.Dict[str, str] + + +def test_jsondc_schemaize(): + lt = TypeEngine.to_literal_type(Foo) + pt = TypeEngine.guess_python_type(lt) + + # When postponed annotations are enabled, dataclass_json will not work and we'll end up with a + # schemaless generic. + # This test basically tests the broken behavior. Remove this test if + # https://github.com/lovasoa/marshmallow_dataclass/issues/13 is ever fixed. + assert pt is dict From 64cc6e1c449440721bbb5f541e986a2e9d0f8d69 Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Mon, 13 Dec 2021 18:30:34 -0800 Subject: [PATCH 034/128] Complex dataclass unit tests (#773) Signed-off-by: Ketan Umare Signed-off-by: Yee Hing Tong Signed-off-by: maximsmol --- .../unit/core/test_complex_nesting.py | 237 ++++++++++++++++++ 1 file changed, 237 insertions(+) create mode 100644 tests/flytekit/unit/core/test_complex_nesting.py diff --git a/tests/flytekit/unit/core/test_complex_nesting.py b/tests/flytekit/unit/core/test_complex_nesting.py new file mode 100644 index 0000000000..13bb108f7c --- /dev/null +++ b/tests/flytekit/unit/core/test_complex_nesting.py @@ -0,0 +1,237 @@ +import os +import tempfile +from dataclasses import dataclass +from typing import List + +import pytest +from dataclasses_json import dataclass_json + +from flytekit.core.context_manager import ExecutionState, FlyteContextManager, Image, ImageConfig, SerializationSettings +from flytekit.core.dynamic_workflow_task import dynamic +from flytekit.core.type_engine import TypeEngine +from flytekit.types.directory import FlyteDirectory +from flytekit.types.file import FlyteFile + + +@dataclass_json +@dataclass +class MyProxyConfiguration: + # File and directory paths kept as 'str' so Flyte doesn't manage these static resources + splat_data_dir: str + apriori_file: str + + +@dataclass_json +@dataclass +class MyProxyParameters: + id: str + job_i_step: int + + +@dataclass_json +@dataclass +class MyAprioriConfiguration: + static_data_dir: FlyteDirectory + external_data_dir: FlyteDirectory + + +@dataclass_json +@dataclass +class MyInput: + main_product: FlyteFile + apriori_config: MyAprioriConfiguration + proxy_config: MyProxyConfiguration + proxy_params: MyProxyParameters + + +@pytest.fixture +def folders_and_files_setup(): + tmp_dir = tempfile.TemporaryDirectory() + fd, path = tempfile.mkstemp(dir=tmp_dir.name) + tmp_dir_static_data = tempfile.TemporaryDirectory() + tmp_dir_external_data = tempfile.TemporaryDirectory() + + try: + with os.fdopen(fd, "w") as tmp: + tmp.write("Hello world") + yield path, tmp_dir_static_data.name, tmp_dir_external_data.name + finally: + tmp_dir.cleanup() + tmp_dir_static_data.cleanup() + tmp_dir_external_data.cleanup() + + +@pytest.fixture +def two_sample_inputs(folders_and_files_setup): + (file_path, static_data_path, external_data_path) = folders_and_files_setup + + main_product = FlyteFile(file_path) + apriori = MyAprioriConfiguration( + static_data_dir=FlyteDirectory(static_data_path), + external_data_dir=FlyteDirectory(external_data_path), + ) + proxy_c = MyProxyConfiguration(splat_data_dir="/tmp/proxy_splat", apriori_file="/opt/config/a_file") + proxy_p = MyProxyParameters(id="pp_id", job_i_step=1) + + my_input = MyInput( + main_product=main_product, + apriori_config=apriori, + proxy_config=proxy_c, + proxy_params=proxy_p, + ) + + my_input_2 = MyInput( + main_product=main_product, + apriori_config=apriori, + proxy_config=proxy_c, + proxy_params=proxy_p, + ) + + yield my_input, my_input_2 + + +def test_dataclass_complex_transform(two_sample_inputs): + my_input = two_sample_inputs[0] + my_input_2 = two_sample_inputs[1] + + ctx = FlyteContextManager.current_context() + literal_type = TypeEngine.to_literal_type(MyInput) + first_literal = TypeEngine.to_literal(ctx, my_input, MyInput, literal_type) + assert first_literal.scalar.generic["apriori_config"] is not None + + converted_back_1 = TypeEngine.to_python_value(ctx, first_literal, MyInput) + assert converted_back_1.apriori_config is not None + + second_literal = TypeEngine.to_literal(ctx, converted_back_1, MyInput, literal_type) + assert second_literal.scalar.generic["apriori_config"] is not None + + converted_back_2 = TypeEngine.to_python_value(ctx, second_literal, MyInput) + assert converted_back_2.apriori_config is not None + + input_list = [my_input, my_input_2] + input_list_type = TypeEngine.to_literal_type(List[MyInput]) + literal_list = TypeEngine.to_literal(ctx, input_list, List[MyInput], input_list_type) + assert literal_list.collection.literals[0].scalar.generic["apriori_config"] is not None + assert literal_list.collection.literals[1].scalar.generic["apriori_config"] is not None + + +def test_two(two_sample_inputs): + my_input = two_sample_inputs[0] + my_input_2 = two_sample_inputs[1] + + @dynamic + def dt1(a: List[MyInput]) -> List[FlyteFile]: + x = [] + for aa in a: + x.append(aa.main_product) + return x + + with FlyteContextManager.with_context( + FlyteContextManager.current_context().with_serialization_settings( + SerializationSettings( + project="test_proj", + domain="test_domain", + version="abc", + image_config=ImageConfig(Image(name="name", fqn="image", tag="name")), + env={}, + ) + ) + ) as ctx: + with FlyteContextManager.with_context( + ctx.with_execution_state( + ctx.execution_state.with_params( + mode=ExecutionState.Mode.TASK_EXECUTION, + additional_context={ + "dynamic_addl_distro": "s3://my-s3-bucket/fast/123", + "dynamic_dest_dir": "/User/flyte/workflows", + }, + ) + ) + ) as ctx: + input_literal_map = TypeEngine.dict_to_literal_map( + ctx, d={"a": [my_input, my_input_2]}, guessed_python_types={"a": List[MyInput]} + ) + dynamic_job_spec = dt1.dispatch_execute(ctx, input_literal_map) + assert len(dynamic_job_spec.literals["o0"].collection.literals) == 2 + + +def test_str_input(folders_and_files_setup): + proxy_c = MyProxyConfiguration(splat_data_dir="/tmp/proxy_splat", apriori_file="/opt/config/a_file") + proxy_p = MyProxyParameters(id="pp_id", job_i_step=1) + + # Intentionally passing in the wrong type + my_input = MyInput( + main_product=folders_and_files_setup[0], # noqa + apriori_config=MyAprioriConfiguration( + static_data_dir=FlyteDirectory("gs://my-bucket/one"), + external_data_dir=FlyteDirectory("gs://my-bucket/two"), + ), + proxy_config=proxy_c, + proxy_params=proxy_p, + ) + ctx = FlyteContextManager.current_context() + literal_type = TypeEngine.to_literal_type(MyInput) + first_literal = TypeEngine.to_literal(ctx, my_input, MyInput, literal_type) + assert first_literal.scalar.generic is not None + + +def test_dc_dyn_directory(folders_and_files_setup): + proxy_c = MyProxyConfiguration(splat_data_dir="/tmp/proxy_splat", apriori_file="/opt/config/a_file") + proxy_p = MyProxyParameters(id="pp_id", job_i_step=1) + + my_input_gcs = MyInput( + main_product=FlyteFile(folders_and_files_setup[0]), + apriori_config=MyAprioriConfiguration( + static_data_dir=FlyteDirectory("gs://my-bucket/one"), + external_data_dir=FlyteDirectory("gs://my-bucket/two"), + ), + proxy_config=proxy_c, + proxy_params=proxy_p, + ) + + my_input_gcs_2 = MyInput( + main_product=FlyteFile(folders_and_files_setup[0]), + apriori_config=MyAprioriConfiguration( + static_data_dir=FlyteDirectory("gs://my-bucket/three"), + external_data_dir=FlyteDirectory("gs://my-bucket/four"), + ), + proxy_config=proxy_c, + proxy_params=proxy_p, + ) + + @dynamic + def dt1(a: List[MyInput]) -> List[FlyteDirectory]: + x = [] + for aa in a: + x.append(aa.apriori_config.external_data_dir) + + return x + + with FlyteContextManager.with_context( + FlyteContextManager.current_context().with_serialization_settings( + SerializationSettings( + project="test_proj", + domain="test_domain", + version="abc", + image_config=ImageConfig(Image(name="name", fqn="image", tag="name")), + env={}, + ) + ) + ) as ctx: + with FlyteContextManager.with_context( + ctx.with_execution_state( + ctx.execution_state.with_params( + mode=ExecutionState.Mode.TASK_EXECUTION, + additional_context={ + "dynamic_addl_distro": "s3://my-s3-bucket/fast/123", + "dynamic_dest_dir": "/User/flyte/workflows", + }, + ) + ) + ) as ctx: + input_literal_map = TypeEngine.dict_to_literal_map( + ctx, d={"a": [my_input_gcs, my_input_gcs_2]}, guessed_python_types={"a": List[MyInput]} + ) + dynamic_job_spec = dt1.dispatch_execute(ctx, input_literal_map) + assert dynamic_job_spec.literals["o0"].collection.literals[0].scalar.blob.uri == "gs://my-bucket/two" + assert dynamic_job_spec.literals["o0"].collection.literals[1].scalar.blob.uri == "gs://my-bucket/four" From 41ce10d44bb9cf8176cb3d45df51605616ca9c2c Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Wed, 15 Dec 2021 03:18:42 +0800 Subject: [PATCH 035/128] remote_source lost on serialization of @dataclass_json with FlyteFile (#774) * remote_source lost on serialization of @dataclass_json with FlyteFile Signed-off-by: Kevin Su * updated tests Signed-off-by: Kevin Su * updated tests Signed-off-by: Kevin Su Signed-off-by: maximsmol --- flytekit/core/type_engine.py | 24 +++++++++++++++----- tests/flytekit/unit/core/test_flyte_file.py | 25 +++++++++++++++++++++ tests/flytekit/unit/core/test_type_hints.py | 25 ++++++++++++++++++--- 3 files changed, 66 insertions(+), 8 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 178b550136..21c0749fbb 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -311,12 +311,26 @@ def _serialize_flyte_type(self, python_val: T, python_type: Type[T]): for f in dataclasses.fields(python_type): v = python_val.__getattribute__(f.name) - if inspect.isclass(f.type) and ( - issubclass(f.type, FlyteSchema) or issubclass(f.type, FlyteFile) or issubclass(f.type, FlyteDirectory) + field_type = f.type + if inspect.isclass(field_type) and ( + issubclass(field_type, FlyteSchema) + or issubclass(field_type, FlyteFile) + or issubclass(field_type, FlyteDirectory) ): - TypeEngine.to_literal(FlyteContext.current_context(), v, f.type, None) - elif dataclasses.is_dataclass(f.type): - self._serialize_flyte_type(v, f.type) + lv = TypeEngine.to_literal(FlyteContext.current_context(), v, field_type, None) + # dataclass_json package will extract the "path" from FlyteFile, FlyteDirectory, and write it to a + # JSON which will be stored in IDL. The path here should always be a remote path, but sometimes the + # path in FlyteFile and FlyteDirectory could be a local path. Therefore, reset the python value here, + # so that dataclass_json can always get a remote path. + # In other words, the file transformer has special code that handles the fact that if remote_source is + # set, then the real uri in the literal should be the remote source, not the path (which may be an + # auto-generated random local path). To be sure we're writing the right path to the json, use the uri + # as determined by the transformer. + if issubclass(field_type, FlyteFile) or issubclass(field_type, FlyteDirectory): + python_val.__setattr__(f.name, field_type(path=lv.scalar.blob.uri)) + + elif dataclasses.is_dataclass(field_type): + self._serialize_flyte_type(v, field_type) def _deserialize_flyte_type(self, python_val: T, expected_python_type: Type) -> T: from flytekit.types.directory.types import FlyteDirectory, FlyteDirToMultipartBlobTransformer diff --git a/tests/flytekit/unit/core/test_flyte_file.py b/tests/flytekit/unit/core/test_flyte_file.py index 3a34445f08..05d8156ece 100644 --- a/tests/flytekit/unit/core/test_flyte_file.py +++ b/tests/flytekit/unit/core/test_flyte_file.py @@ -407,3 +407,28 @@ def test_file_guess(): fft = transformer.guess_python_type(lt) assert issubclass(fft, FlyteFile) assert fft.extension() == "" + + +def test_flyte_file_in_dyn(): + @task + def t1(path: str) -> FlyteFile: + return FlyteFile(path) + + @dynamic + def dyn(fs: FlyteFile): + t2(ff=fs) + + @task + def t2(ff: FlyteFile) -> os.PathLike: + assert ff.remote_source == "s3://somewhere" + assert "/tmp/flyte/" in ff.path + + return ff.path + + @workflow + def wf(path: str) -> os.PathLike: + n1 = t1(path=path) + dyn(fs=n1) + return t2(ff=n1) + + assert "/tmp/flyte/" in wf(path="s3://somewhere").path diff --git a/tests/flytekit/unit/core/test_type_hints.py b/tests/flytekit/unit/core/test_type_hints.py index 131bd7cc09..2f4a4e3068 100644 --- a/tests/flytekit/unit/core/test_type_hints.py +++ b/tests/flytekit/unit/core/test_type_hints.py @@ -372,16 +372,35 @@ def t1(path: str) -> FileStruct: fs = FileStruct(a=file, b=InnerFileStruct(a=file, b=PNGImageFile(path))) return fs + @dynamic + def dyn(fs: FileStruct): + t2(fs=fs) + t3(fs=fs) + @task def t2(fs: FileStruct) -> os.PathLike: + assert fs.a.remote_source == "s3://somewhere" + assert fs.b.a.remote_source == "s3://somewhere" + assert fs.b.b.remote_source == "s3://somewhere" + assert "/tmp/flyte/" in fs.a.path + assert "/tmp/flyte/" in fs.b.a.path + assert "/tmp/flyte/" in fs.b.b.path + return fs.a.path + @task + def t3(fs: FileStruct) -> FlyteFile: + return fs.a + @workflow - def wf(path: str) -> os.PathLike: + def wf(path: str) -> (os.PathLike, FlyteFile): n1 = t1(path=path) - return t2(fs=n1) + dyn(fs=n1) + return t2(fs=n1), t3(fs=n1) - assert "/tmp/flyte/" in wf(path="s3://somewhere").path + assert "/tmp/flyte/" in wf(path="s3://somewhere")[0].path + assert "/tmp/flyte/" in wf(path="s3://somewhere")[1].path + assert "s3://somewhere" == wf(path="s3://somewhere")[1].remote_source def test_flyte_directory_in_dataclass(): From 9a4c582ec255cae0f8c371f90e2fe9f2dccc681f Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Tue, 14 Dec 2021 14:19:45 -0800 Subject: [PATCH 036/128] Single-task execution FlyteRemote sync (#778) Signed-off-by: Yee Hing Tong Signed-off-by: maximsmol --- flytekit/remote/remote.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/flytekit/remote/remote.py b/flytekit/remote/remote.py index 60ae7bbf3e..2178f148d3 100644 --- a/flytekit/remote/remote.py +++ b/flytekit/remote/remote.py @@ -1117,6 +1117,14 @@ def sync_node_execution( """ # For single task execution - the metadata spec node id is missing. In these cases, revert to regular node id node_id = execution.metadata.spec_node_id + # This case supports single-task execution compiled workflows. + if node_id and node_id not in node_mapping and execution.id.node_id in node_mapping: + node_id = execution.id.node_id + remote_logger.debug( + f"Using node execution ID {node_id} instead of spec node id " + f"{execution.metadata.spec_node_id}, single-task execution likely." + ) + # This case supports single-task execution compiled workflows with older versions of admin/propeller if not node_id: node_id = execution.id.node_id remote_logger.debug(f"No metadata spec_node_id found, using {node_id}") From 7ec8beb7edd2bce702cda64884bfc3610d31ab4b Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Tue, 14 Dec 2021 14:43:59 -0800 Subject: [PATCH 037/128] Logging updates (#775) Signed-off-by: Yee Hing Tong Signed-off-by: maximsmol --- flytekit/bin/entrypoint.py | 59 ++++++++++++++++---------------- flytekit/configuration/images.py | 3 +- flytekit/loggers.py | 57 +++++++++++++++++++++++------- 3 files changed, 76 insertions(+), 43 deletions(-) diff --git a/flytekit/bin/entrypoint.py b/flytekit/bin/entrypoint.py index 6153ed6fa9..1888376e50 100644 --- a/flytekit/bin/entrypoint.py +++ b/flytekit/bin/entrypoint.py @@ -1,7 +1,7 @@ import contextlib import datetime as _datetime import importlib as _importlib -import logging as _logging +import logging as python_logging import os as _os import pathlib import random as _random @@ -37,6 +37,7 @@ from flytekit.interfaces import random as _flyte_random from flytekit.interfaces.data import data_proxy as _data_proxy from flytekit.interfaces.stats.taggable import get_stats as _get_stats +from flytekit.loggers import entrypoint_logger as logger from flytekit.models import dynamic_job as _dynamic_job from flytekit.models import literals as _literal_models from flytekit.models.core import errors as _error_models @@ -95,6 +96,7 @@ def _dispatch_execute( c: OR if an unhandled exception is retrieved - record it as an errors.pb """ output_file_dict = {} + logger.debug(f"Starting _dispatch_execute for {task_def.name}") try: # Step1 local_inputs_file = _os.path.join(ctx.execution_state.working_dir, "inputs.pb") @@ -108,14 +110,14 @@ def _dispatch_execute( outputs = _scoped_exceptions.system_entry_point(task_def.dispatch_execute)(ctx, idl_input_literals) # Step3a if isinstance(outputs, VoidPromise): - _logging.getLogger().warning("Task produces no outputs") + logger.warning("Task produces no outputs") output_file_dict = {_constants.OUTPUT_FILE_NAME: _literal_models.LiteralMap(literals={})} elif isinstance(outputs, _literal_models.LiteralMap): output_file_dict = {_constants.OUTPUT_FILE_NAME: outputs} elif isinstance(outputs, _dynamic_job.DynamicJobSpec): output_file_dict = {_constants.FUTURES_FILE_NAME: outputs} else: - _logging.getLogger().error(f"SystemError: received unknown outputs from task {outputs}") + logger.error(f"SystemError: received unknown outputs from task {outputs}") output_file_dict[_constants.ERROR_FILE_NAME] = _error_models.ErrorDocument( _error_models.ContainerError( "UNKNOWN_OUTPUT", @@ -128,30 +130,30 @@ def _dispatch_execute( # Handle user-scoped errors except _scoped_exceptions.FlyteScopedUserException as e: if isinstance(e.value, IgnoreOutputs): - _logging.warning(f"User-scoped IgnoreOutputs received! Outputs.pb will not be uploaded. reason {e}!!") + logger.warning(f"User-scoped IgnoreOutputs received! Outputs.pb will not be uploaded. reason {e}!!") return output_file_dict[_constants.ERROR_FILE_NAME] = _error_models.ErrorDocument( _error_models.ContainerError( e.error_code, e.verbose_message, e.kind, _execution_models.ExecutionError.ErrorKind.USER ) ) - _logging.error("!! Begin User Error Captured by Flyte !!") - _logging.error(e.verbose_message) - _logging.error("!! End Error Captured by Flyte !!") + logger.error("!! Begin User Error Captured by Flyte !!") + logger.error(e.verbose_message) + logger.error("!! End Error Captured by Flyte !!") # Handle system-scoped errors except _scoped_exceptions.FlyteScopedSystemException as e: if isinstance(e.value, IgnoreOutputs): - _logging.warning(f"System-scoped IgnoreOutputs received! Outputs.pb will not be uploaded. reason {e}!!") + logger.warning(f"System-scoped IgnoreOutputs received! Outputs.pb will not be uploaded. reason {e}!!") return output_file_dict[_constants.ERROR_FILE_NAME] = _error_models.ErrorDocument( _error_models.ContainerError( e.error_code, e.verbose_message, e.kind, _execution_models.ExecutionError.ErrorKind.SYSTEM ) ) - _logging.error("!! Begin System Error Captured by Flyte !!") - _logging.error(e.verbose_message) - _logging.error("!! End Error Captured by Flyte !!") + logger.error("!! Begin System Error Captured by Flyte !!") + logger.error(e.verbose_message) + logger.error("!! End Error Captured by Flyte !!") # Interpret all other exceptions (some of which may be caused by the code in the try block outside of # dispatch_execute) as recoverable system exceptions. @@ -166,16 +168,17 @@ def _dispatch_execute( _execution_models.ExecutionError.ErrorKind.SYSTEM, ) ) - _logging.error(f"Exception when executing task {task_def.name or task_def.id.name}, reason {str(e)}") - _logging.error("!! Begin Unknown System Error Captured by Flyte !!") - _logging.error(exc_str) - _logging.error("!! End Error Captured by Flyte !!") + logger.error(f"Exception when executing task {task_def.name or task_def.id.name}, reason {str(e)}") + logger.error("!! Begin Unknown System Error Captured by Flyte !!") + logger.error(exc_str) + logger.error("!! End Error Captured by Flyte !!") for k, v in output_file_dict.items(): _common_utils.write_proto_to_file(v.to_flyte_idl(), _os.path.join(ctx.execution_state.engine_dir, k)) ctx.file_access.put_data(ctx.execution_state.engine_dir, output_prefix, is_multipart=True) - _logging.info(f"Engine folder written successfully to the output prefix {output_prefix}") + logger.info(f"Engine folder written successfully to the output prefix {output_prefix}") + logger.debug("Finished _dispatch_execute") @contextlib.contextmanager @@ -184,14 +187,11 @@ def setup_execution( dynamic_addl_distro: str = None, dynamic_dest_dir: str = None, ): - log_level = _internal_config.LOGGING_LEVEL.get() or _sdk_config.LOGGING_LEVEL.get() - _logging.getLogger().setLevel(log_level) - ctx = FlyteContextManager.current_context() # Create directories user_workspace_dir = ctx.file_access.get_random_local_directory() - _click.echo(f"Using user directory {user_workspace_dir}") + logger.info(f"Using user directory {user_workspace_dir}") pathlib.Path(user_workspace_dir).mkdir(parents=True, exist_ok=True) from flytekit import __version__ as _api_version @@ -219,7 +219,7 @@ def setup_execution( "api_version": _api_version, }, ), - logging=_logging, + logging=python_logging, tmp_dir=user_workspace_dir, ) @@ -231,7 +231,7 @@ def setup_execution( raw_output_prefix=raw_output_data_prefix, ) except TypeError: # would be thrown from DataPersistencePlugins.find_plugin - _logging.error(f"No data plugin found for raw output prefix {raw_output_data_prefix}") + logger.error(f"No data plugin found for raw output prefix {raw_output_data_prefix}") raise else: raise Exception("No raw output prefix detected. Please upgrade your version of Propeller to 0.4.0 or later.") @@ -280,7 +280,6 @@ def _handle_annotated_task( """ Entrypoint for all PythonTask extensions """ - _click.echo("Running native-typed task") _dispatch_execute(ctx, task_def, inputs, output_prefix) @@ -366,7 +365,7 @@ def _execute_task( # Use the resolver to load the actual task object _task_def = resolver_obj.load_task(loader_args=resolver_args) if test: - _click.echo( + logger.info( f"Test detected, returning. Args were {inputs} {output_prefix} {raw_output_data_prefix} {resolver} {resolver_args}" ) return @@ -401,7 +400,7 @@ def _execute_map_task( output_prefix = _os.path.join(output_prefix, str(task_index)) if test: - _click.echo( + logger.info( f"Test detected, returning. Inputs: {inputs} Computed task index: {task_index} " f"New output prefix: {output_prefix} Raw output path: {raw_output_data_prefix} " f"Resolver and args: {resolver} {resolver_args}" @@ -443,7 +442,9 @@ def execute_task_cmd( resolver, resolver_args, ): - _click.echo(_utils.get_version_message()) + logger.info(_utils.get_version_message()) + # We get weird errors if there are no click echo messages at all, so emit an empty string so that unit tests pass. + _click.echo("") # Backwards compatibility - if Propeller hasn't filled this in, then it'll come through here as the original # template string, so let's explicitly set it to None so that the downstream functions will know to fall back # to the original shard formatter/prefix config. @@ -455,10 +456,10 @@ def execute_task_cmd( # The addition of a new top-level command seemed out of scope at the time of this writing to pursue given how # pervasive this top level command already (plugins mostly). if not resolver: - _click.echo("No resolver found, assuming legacy API task...") + logger.info("No resolver found, assuming legacy API task...") _legacy_execute_task(task_module, task_name, inputs, output_prefix, raw_output_data_prefix, test) else: - _click.echo(f"Attempting to run with {resolver}...") + logger.debug(f"Running task execution with resolver {resolver}...") _execute_task( inputs, output_prefix, @@ -527,7 +528,7 @@ def map_execute_task_cmd( resolver, resolver_args, ): - _click.echo(_utils.get_version_message()) + logger.info(_utils.get_version_message()) _execute_map_task( inputs, diff --git a/flytekit/configuration/images.py b/flytekit/configuration/images.py index 10544071d6..092ccaaff7 100644 --- a/flytekit/configuration/images.py +++ b/flytekit/configuration/images.py @@ -2,6 +2,7 @@ import typing from flytekit.configuration import common as _config_common +from flytekit.loggers import logger def get_specified_images() -> typing.Dict[str, str]: @@ -21,7 +22,7 @@ def get_specified_images() -> typing.Dict[str, str]: try: image_names = _config_common.CONFIGURATION_SINGLETON.config.options("images") except configparser.NoSectionError: - print("No images specified, will use the default image") + logger.info("No images specified, will use the default image") image_names = None if image_names: for i in image_names: diff --git a/flytekit/loggers.py b/flytekit/loggers.py index 4527a659e1..bc3e243883 100644 --- a/flytekit/loggers.py +++ b/flytekit/loggers.py @@ -1,29 +1,60 @@ -import logging as _logging -import os as _os +import logging +import os from pythonjsonlogger import jsonlogger -logger = _logging.getLogger("flytekit") -# Always set the root logger to debug until we can add more user based controls -logger.setLevel(_logging.WARNING) +# Note: +# The environment variable controls exposed to affect the individual loggers should be considered to be beta. +# The ux/api may change in the future. +# At time of writing, the code was written to preserve existing default behavior +# For now, assume this is the environment variable whose usage will remain unchanged and controls output for all +# loggers defined in this file. +LOGGING_ENV_VAR = "FLYTE_SDK_LOGGING_LEVEL" + +# By default, the root flytekit logger to debug so everything is logged, but enable fine-tuning +logger = logging.getLogger("flytekit") +# Root logger control +flytekit_root_env_var = f"{LOGGING_ENV_VAR}_ROOT" +if os.getenv(flytekit_root_env_var) is not None: + logger.setLevel(int(os.getenv(flytekit_root_env_var))) +else: + logger.setLevel(logging.DEBUG) + +# Stop propagation so that configuration is isolated to this file (so that it doesn't matter what the +# global Python root logger is set to). +logger.propagate = False # Child loggers -auth_logger = logger.getChild("auth") -cli_logger = logger.getChild("cli") -remote_logger = logger.getChild("remote") +child_loggers = { + "auth": logger.getChild("auth"), + "cli": logger.getChild("cli"), + "remote": logger.getChild("remote"), + "entrypoint": logger.getChild("entrypoint"), +} +auth_logger = child_loggers["auth"] +cli_logger = child_loggers["cli"] +remote_logger = child_loggers["remote"] +entrypoint_logger = child_loggers["entrypoint"] -# create console handler and set level to debug -ch = _logging.StreamHandler() +# create console handler +ch = logging.StreamHandler() # Don't want to import the configuration library since that will cause all sorts of circular imports, let's # just use the environment variable if it's defined. Decide in the future when we implement better controls # if we should control with the channel or with the logger level. -logging_env_var = "FLYTE_SDK_LOGGING_LEVEL" -level_from_env = _os.getenv(logging_env_var) +# The handler log level controls whether log statements will actually print to the screen +level_from_env = os.getenv(LOGGING_ENV_VAR) if level_from_env is not None: ch.setLevel(int(level_from_env)) else: - ch.setLevel(_logging.WARNING) + ch.setLevel(logging.WARNING) + +# Consider this API to be beta +for log_name, child_logger in child_loggers.items(): + env_var = f"{LOGGING_ENV_VAR}_{log_name.upper()}" + level_from_env = os.getenv(env_var) + if level_from_env is not None: + child_logger.setLevel(int(level_from_env)) # create formatter formatter = jsonlogger.JsonFormatter(fmt="%(asctime)s %(name)s %(levelname)s %(message)s") From 68dd191f478bdbbdedc44a2c53140464ae24dfe1 Mon Sep 17 00:00:00 2001 From: Snyk bot Date: Wed, 15 Dec 2021 18:50:03 +0000 Subject: [PATCH 038/128] fix: doc-requirements.txt to reduce vulnerabilities (#779) The following vulnerabilities are fixed by pinning transitive dependencies: - https://snyk.io/vuln/SNYK-PYTHON-LXML-2316995 Signed-off-by: maximsmol --- doc-requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/doc-requirements.txt b/doc-requirements.txt index 1e766d5f31..715b03e2e3 100644 --- a/doc-requirements.txt +++ b/doc-requirements.txt @@ -165,7 +165,7 @@ keyring==23.2.1 # via flytekit lazy-object-proxy==1.6.0 # via astroid -lxml==4.6.4 +lxml==4.6.5 # via sphinx-material markupsafe==2.0.1 # via jinja2 From 877038a7321793f873b77ac2722f1059cfb9a626 Mon Sep 17 00:00:00 2001 From: Dan Rammer Date: Wed, 15 Dec 2021 16:54:48 -0600 Subject: [PATCH 039/128] Add cache_serialize parameter to tasks (#673) * added cache_serialize parameter for basic tasks Signed-off-by: Daniel Rammer * fixed typo Signed-off-by: Daniel Rammer * updated flyteidl version Signed-off-by: Daniel Rammer * changed flyteidl requirements everywhere Signed-off-by: Daniel Rammer * remove flyteidl version requires in setup.py so we can use develop Signed-off-by: Daniel Rammer * removed flyteidl git repos from a variety of requirements packages Signed-off-by: Daniel Rammer * updated variable discovery_serializable to cache_serializable Signed-off-by: Daniel Rammer * updated requirements Signed-off-by: Daniel Rammer * fixed TaskMetadata _cache_serializable variable name Signed-off-by: Daniel Rammer * propgating cache_serialize parameter through to tasks Signed-off-by: Daniel Rammer * added cache_serializable to SdkRawContainerTask Signed-off-by: Daniel Rammer * fixing cache_serializable variable propogation issues Signed-off-by: Daniel Rammer * added documentation Signed-off-by: Daniel Rammer * added unit tests for cache_serialize metadata Signed-off-by: Daniel Rammer * linter added spaces in unit tests Signed-off-by: Daniel Rammer Signed-off-by: maximsmol --- dev-requirements.txt | 42 +++++------ doc-requirements.txt | 72 +++++++++---------- flytekit/common/tasks/hive_task.py | 3 + flytekit/common/tasks/presto_task.py | 3 + flytekit/common/tasks/raw_container.py | 3 + flytekit/common/tasks/sdk_dynamic.py | 3 + flytekit/common/tasks/sdk_runnable.py | 3 + flytekit/common/tasks/sidecar_task.py | 5 ++ flytekit/common/tasks/spark_task.py | 3 + flytekit/contrib/sensors/task.py | 1 + flytekit/core/base_task.py | 5 ++ flytekit/core/task.py | 6 ++ flytekit/models/task.py | 15 ++++ flytekit/sdk/tasks.py | 39 ++++++++++ requirements-spark2.txt | 56 +++++++-------- requirements.txt | 56 +++++++-------- tests/flytekit/common/parameterizers.py | 4 +- .../workflows/requirements.txt | 49 ++++++++----- .../common_tests/tasks/test_sdk_runnable.py | 1 + .../common_tests/test_workflow_promote.py | 1 + .../unit/core/test_python_function_task.py | 34 +++++++-- tests/flytekit/unit/models/test_tasks.py | 2 + .../unit/models/test_workflow_closure.py | 1 + 23 files changed, 270 insertions(+), 137 deletions(-) diff --git a/dev-requirements.txt b/dev-requirements.txt index a1b1fa3827..4ba14b6a8c 100644 --- a/dev-requirements.txt +++ b/dev-requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.8 +# This file is autogenerated by pip-compile with python 3.9 # To update, run: # # make dev-requirements.txt @@ -44,7 +44,7 @@ chardet==4.0.0 # via # -c requirements.txt # binaryornot -charset-normalizer==2.0.7 +charset-normalizer==2.0.8 # via # -c requirements.txt # requests @@ -67,13 +67,13 @@ cookiecutter==1.7.3 # via # -c requirements.txt # flytekit -coverage[toml]==6.1.1 +coverage[toml]==6.2 # via -r dev-requirements.in croniter==1.0.15 # via # -c requirements.txt # flytekit -cryptography==35.0.0 +cryptography==36.0.0 # via # -c requirements.txt # paramiko @@ -90,7 +90,7 @@ deprecated==1.2.13 # via # -c requirements.txt # flytekit -diskcache==5.2.1 +diskcache==5.3.0 # via # -c requirements.txt # flytekit @@ -112,21 +112,21 @@ dockerpty==0.4.1 # via docker-compose docopt==0.6.2 # via docker-compose -docstring-parser==0.12 +docstring-parser==0.13 # via # -c requirements.txt # flytekit -filelock==3.3.2 +filelock==3.4.0 # via virtualenv -flyteidl==0.21.8 +flyteidl==0.21.11 # via # -c requirements.txt # flytekit -grpcio==1.41.1 +grpcio==1.42.0 # via # -c requirements.txt # flytekit -identify==2.3.5 +identify==2.4.0 # via pre-commit idna==3.3 # via @@ -159,7 +159,7 @@ jsonschema==3.2.0 # via # -c requirements.txt # docker-compose -keyring==23.2.1 +keyring==23.4.0 # via # -c requirements.txt # flytekit @@ -167,7 +167,7 @@ markupsafe==2.0.1 # via # -c requirements.txt # jinja2 -marshmallow==3.14.0 +marshmallow==3.14.1 # via # -c requirements.txt # dataclasses-json @@ -201,7 +201,7 @@ numpy==1.21.4 # -c requirements.txt # pandas # pyarrow -packaging==21.2 +packaging==21.3 # via # -c requirements.txt # pytest @@ -209,7 +209,7 @@ pandas==1.3.4 # via # -c requirements.txt # flytekit -paramiko==2.8.0 +paramiko==2.8.1 # via # -c requirements.txt # docker @@ -223,7 +223,7 @@ poyo==0.5.0 # via # -c requirements.txt # cookiecutter -pre-commit==2.15.0 +pre-commit==2.16.0 # via -r dev-requirements.in protobuf==3.19.1 # via @@ -235,7 +235,7 @@ py==1.11.0 # -c requirements.txt # pytest # retry -pyarrow==6.0.0 +pyarrow==6.0.1 # via # -c requirements.txt # flytekit @@ -247,7 +247,7 @@ pynacl==1.4.0 # via # -c requirements.txt # paramiko -pyparsing==2.4.7 +pyparsing==3.0.6 # via # -c requirements.txt # packaging @@ -271,7 +271,7 @@ python-dateutil==2.8.1 # croniter # flytekit # pandas -python-dotenv==0.19.1 +python-dotenv==0.19.2 # via docker-compose python-json-logger==2.0.2 # via @@ -295,7 +295,7 @@ pyyaml==5.4.1 # -c requirements.txt # docker-compose # pre-commit -regex==2021.11.9 +regex==2021.11.10 # via # -c requirements.txt # docker-image-py @@ -307,7 +307,7 @@ requests==2.26.0 # docker-compose # flytekit # responses -responses==0.15.0 +responses==0.16.0 # via # -c requirements.txt # flytekit @@ -356,7 +356,7 @@ tomli==1.2.2 # via # -c requirements.txt # coverage -typing-extensions==3.10.0.2 +typing-extensions==4.0.1 # via # -c requirements.txt # mypy diff --git a/doc-requirements.txt b/doc-requirements.txt index 715b03e2e3..1913e64cf2 100644 --- a/doc-requirements.txt +++ b/doc-requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.8 +# This file is autogenerated by pip-compile with python 3.9 # To update, run: # # make doc-requirements.txt @@ -12,7 +12,7 @@ ansiwrap==0.8.4 # via papermill arrow==1.2.1 # via jinja2-time -astroid==2.8.4 +astroid==2.9.0 # via sphinx-autoapi attrs==21.2.0 # via jsonschema @@ -29,13 +29,13 @@ beautifulsoup4==4.10.0 # sphinx-material binaryornot==0.4.4 # via cookiecutter -black==21.10b0 +black==21.11b1 # via papermill bleach==4.1.0 # via nbconvert -boto3==1.20.2 +boto3==1.20.17 # via sagemaker-training -botocore==1.23.2 +botocore==1.23.17 # via # boto3 # s3transfer @@ -48,7 +48,7 @@ cffi==1.15.0 # pynacl chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.7 +charset-normalizer==2.0.8 # via requests checksumdir==1.2.0 # via flytekit @@ -65,7 +65,7 @@ cookiecutter==1.7.3 # via flytekit croniter==1.0.15 # via flytekit -cryptography==35.0.0 +cryptography==36.0.0 # via # -r doc-requirements.in # paramiko @@ -82,11 +82,11 @@ defusedxml==0.7.1 # via nbconvert deprecated==1.2.13 # via flytekit -diskcache==5.2.1 +diskcache==5.3.0 # via flytekit docker-image-py==0.1.12 # via flytekit -docstring-parser==0.12 +docstring-parser==0.13 # via flytekit docutils==0.17.1 # via sphinx @@ -95,7 +95,7 @@ entrypoints==0.3 # jupyter-client # nbconvert # papermill -flyteidl==0.21.8 +flyteidl==0.21.11 # via flytekit furo @ git+git://github.com/flyteorg/furo@main # via -r doc-requirements.in @@ -103,7 +103,7 @@ gevent==21.8.0 # via sagemaker-training greenlet==1.1.2 # via gevent -grpcio==1.41.1 +grpcio==1.42.0 # via # -r doc-requirements.in # flytekit @@ -115,19 +115,17 @@ imagesize==1.3.0 # via sphinx importlib-metadata==4.8.2 # via keyring -importlib-resources==5.4.0 - # via jsonschema inotify_simple==1.2.1 # via sagemaker-training ipykernel==5.5.6 # via flytekit -ipython==7.29.0 +ipython==7.30.0 # via ipykernel ipython-genutils==0.2.0 # via # ipykernel # nbformat -jedi==0.18.0 +jedi==0.18.1 # via ipython jeepney==0.7.1 # via @@ -148,7 +146,7 @@ jmespath==0.10.0 # botocore jsonschema==4.2.1 # via nbformat -jupyter-client==7.0.6 +jupyter-client==7.1.0 # via # ipykernel # nbclient @@ -161,7 +159,7 @@ jupyterlab-pygments==0.1.2 # via nbconvert k8s-proto==0.0.3 # via flytekit -keyring==23.2.1 +keyring==23.4.0 # via flytekit lazy-object-proxy==1.6.0 # via astroid @@ -169,7 +167,7 @@ lxml==4.6.5 # via sphinx-material markupsafe==2.0.1 # via jinja2 -marshmallow==3.14.0 +marshmallow==3.14.1 # via # dataclasses-json # marshmallow-enum @@ -188,18 +186,18 @@ mypy-extensions==0.4.3 # typing-inspect natsort==8.0.0 # via flytekit -nbclient==0.5.5 +nbclient==0.5.9 # via # nbconvert # papermill -nbconvert==6.2.0 +nbconvert==6.3.0 # via flytekit nbformat==5.1.3 # via # nbclient # nbconvert # papermill -nest-asyncio==1.5.1 +nest-asyncio==1.5.2 # via # jupyter-client # nbclient @@ -210,7 +208,7 @@ numpy==1.21.4 # pyarrow # sagemaker-training # scipy -packaging==21.2 +packaging==21.3 # via # bleach # sphinx @@ -220,9 +218,9 @@ pandocfilters==1.5.0 # via nbconvert papermill==2.3.3 # via flytekit -paramiko==2.8.0 +paramiko==2.8.1 # via sagemaker-training -parso==0.8.2 +parso==0.8.3 # via jedi pathspec==0.9.0 # via black @@ -234,7 +232,7 @@ platformdirs==2.4.0 # via black poyo==0.5.0 # via cookiecutter -prompt-toolkit==3.0.22 +prompt-toolkit==3.0.23 # via ipython protobuf==3.19.1 # via @@ -250,7 +248,7 @@ py==1.11.0 # via retry py4j==0.10.9.2 # via pyspark -pyarrow==6.0.0 +pyarrow==6.0.1 # via flytekit pycparser==2.21 # via cffi @@ -263,7 +261,7 @@ pygments==2.10.0 # sphinx-prompt pynacl==1.4.0 # via paramiko -pyparsing==2.4.7 +pyparsing==3.0.6 # via packaging pyrsistent==0.18.0 # via jsonschema @@ -296,7 +294,7 @@ pyyaml==6.0 # sphinx-autoapi pyzmq==22.3.0 # via jupyter-client -regex==2021.11.9 +regex==2021.11.10 # via # black # docker-image-py @@ -307,7 +305,7 @@ requests==2.26.0 # papermill # responses # sphinx -responses==0.15.0 +responses==0.16.0 # via flytekit retry==0.9.2 # via flytekit @@ -317,7 +315,7 @@ s3transfer==0.5.0 # via boto3 sagemaker-training==3.9.2 # via flytekit -scipy==1.7.2 +scipy==1.7.3 # via sagemaker-training secretstorage==3.3.1 # via keyring @@ -335,13 +333,13 @@ six==1.16.0 # sagemaker-training # sphinx-code-include # thrift -snowballstemmer==2.1.0 +snowballstemmer==2.2.0 # via sphinx sortedcontainers==2.4.0 # via flytekit -soupsieve==2.3 +soupsieve==2.3.1 # via beautifulsoup4 -sphinx==4.2.0 +sphinx==4.3.1 # via # -r doc-requirements.in # furo @@ -361,7 +359,7 @@ sphinx-copybutton==0.4.0 # via -r doc-requirements.in sphinx-fontawesome==0.0.6 # via -r doc-requirements.in -sphinx-gallery==0.10.0 +sphinx-gallery==0.10.1 # via -r doc-requirements.in sphinx-material==0.0.35 # via -r doc-requirements.in @@ -411,7 +409,7 @@ traitlets==5.1.1 # nbclient # nbconvert # nbformat -typing-extensions==3.10.0.2 +typing-extensions==4.0.1 # via # astroid # black @@ -442,9 +440,7 @@ wrapt==1.13.3 # deprecated # flytekit zipp==3.6.0 - # via - # importlib-metadata - # importlib-resources + # via importlib-metadata zope.event==4.5.0 # via gevent zope.interface==5.4.0 diff --git a/flytekit/common/tasks/hive_task.py b/flytekit/common/tasks/hive_task.py index 6be5db70f9..77ae3359a0 100644 --- a/flytekit/common/tasks/hive_task.py +++ b/flytekit/common/tasks/hive_task.py @@ -49,6 +49,7 @@ def __init__( cluster_label, tags, environment, + cache_serializable, ): """ :param task_function: Function container user code. This will be executed via the SDK's engine. @@ -69,6 +70,7 @@ def __init__( :param Text cluster_label: :param list[Text] tags: :param dict[Text, Text] environment: + :param bool cache_serializable: """ self._task_function = task_function super(SdkHiveTask, self).__init__( @@ -89,6 +91,7 @@ def __init__( discoverable, timeout, environment, + cache_serializable, {}, ) self._validate_task_parameters(cluster_label, tags) diff --git a/flytekit/common/tasks/presto_task.py b/flytekit/common/tasks/presto_task.py index 47e198494d..c8f7d300a5 100644 --- a/flytekit/common/tasks/presto_task.py +++ b/flytekit/common/tasks/presto_task.py @@ -35,6 +35,7 @@ def __init__( retries=1, timeout=None, deprecated=None, + cache_serializable=False, ): """ :param Text statement: Presto query specification @@ -49,6 +50,7 @@ def __init__( :param datetime.timedelta timeout: :param Text deprecated: This string can be used to mark the task as deprecated. Consumers of the task will receive deprecation warnings. + :param bool cache_serializable: """ # Set as class fields which are used down below to configure implicit @@ -66,6 +68,7 @@ def __init__( interruptible, discovery_version, deprecated, + cache_serializable, ) presto_query = _presto_models.PrestoQuery( diff --git a/flytekit/common/tasks/raw_container.py b/flytekit/common/tasks/raw_container.py index c290b36205..5168e744f3 100644 --- a/flytekit/common/tasks/raw_container.py +++ b/flytekit/common/tasks/raw_container.py @@ -134,6 +134,7 @@ def __init__( discovery_version: str = None, retries: int = 1, timeout: _datetime.timedelta = None, + cache_serializable: bool = False, ): """ :param inputs: @@ -155,6 +156,7 @@ def __init__( :param discovery_version: :param retries: :param timeout: + :param cache_serializable: :param input_data_dir: This is the directory where data will be downloaded to :param output_data_dir: This is the directory where data will be uploaded from :param metadata_format: Format in which the metadata will be available for the script @@ -183,6 +185,7 @@ def __init__( interruptible, discovery_version, None, + cache_serializable, ) # The interface is defined using the inputs and outputs diff --git a/flytekit/common/tasks/sdk_dynamic.py b/flytekit/common/tasks/sdk_dynamic.py index 2a44111be5..b3bc923601 100644 --- a/flytekit/common/tasks/sdk_dynamic.py +++ b/flytekit/common/tasks/sdk_dynamic.py @@ -320,6 +320,7 @@ def __init__( allowed_failure_ratio, max_concurrency, environment, + cache_serializable, custom, ): """ @@ -342,6 +343,7 @@ def __init__( :param float allowed_failure_ratio: :param int max_concurrency: :param dict[Text, Text] environment: + :param bool cache_serializable: :param dict[Text, T] custom: """ _sdk_runnable.SdkRunnableTask.__init__( @@ -363,6 +365,7 @@ def __init__( discoverable, timeout, environment, + cache_serializable, custom, ) diff --git a/flytekit/common/tasks/sdk_runnable.py b/flytekit/common/tasks/sdk_runnable.py index 1a2ea2af72..39788437dd 100644 --- a/flytekit/common/tasks/sdk_runnable.py +++ b/flytekit/common/tasks/sdk_runnable.py @@ -377,6 +377,7 @@ def __init__( discoverable, timeout, environment, + cache_serializable, custom, ): """ @@ -397,6 +398,7 @@ def __init__( :param bool discoverable: :param datetime.timedelta timeout: :param dict[Text, Text] environment: + :param bool cache_serializable: :param dict[Text, T] custom: """ # Circular dependency @@ -417,6 +419,7 @@ def __init__( interruptible, discovery_version, deprecated, + cache_serializable, ), # TODO: If we end up using SdkRunnableTask for the new code, make sure this is set correctly. _interface.TypedInterface({}, {}), diff --git a/flytekit/common/tasks/sidecar_task.py b/flytekit/common/tasks/sidecar_task.py index d238b166f1..15cb62d760 100644 --- a/flytekit/common/tasks/sidecar_task.py +++ b/flytekit/common/tasks/sidecar_task.py @@ -36,6 +36,7 @@ def __init__( discoverable, timeout, environment, + cache_serializable, pod_spec=None, primary_container_name=None, annotations=None, @@ -72,6 +73,7 @@ def __init__( discoverable, timeout, environment, + cache_serializable, custom=None, ) @@ -180,6 +182,7 @@ def __init__( allowed_failure_ratio, max_concurrency, environment, + cache_serializable, pod_spec=None, primary_container_name=None, annotations=None, @@ -205,6 +208,7 @@ def __init__( :param float allowed_failure_ratio: :param int max_concurrency: :param dict[Text, Text] environment: + :param bool cache_serializable: :param generated_pb2.PodSpec pod_spec: :param Text primary_container_name: :param dict[Text, Text] annotations: @@ -231,6 +235,7 @@ def __init__( discoverable, timeout, environment, + cache_serializable, pod_spec=pod_spec, primary_container_name=primary_container_name, annotations=annotations, diff --git a/flytekit/common/tasks/spark_task.py b/flytekit/common/tasks/spark_task.py index a1dd8251c6..f3f55f211e 100644 --- a/flytekit/common/tasks/spark_task.py +++ b/flytekit/common/tasks/spark_task.py @@ -80,6 +80,7 @@ def __init__( spark_conf, hadoop_conf, environment, + cache_serializable, ): """ :param task_function: Function container user code. This will be executed via the SDK's engine. @@ -93,6 +94,7 @@ def __init__( :param dict[Text,Text] spark_conf: :param dict[Text,Text] hadoop_conf: :param dict[Text,Text] environment: [optional] environment variables to set when executing this task. + :param bool cache_serializable: """ spark_exec_path = _os.path.abspath(_entrypoint.__file__) @@ -125,6 +127,7 @@ def __init__( discoverable, timeout, environment, + cache_serializable, _MessageToDict(self._spark_job.to_flyte_idl()), ) diff --git a/flytekit/contrib/sensors/task.py b/flytekit/contrib/sensors/task.py index 31e8a8d14e..0749fc39dc 100644 --- a/flytekit/contrib/sensors/task.py +++ b/flytekit/contrib/sensors/task.py @@ -116,6 +116,7 @@ def wrapper(fn): custom={}, discovery_version="", discoverable=False, + cache_serializable=False, ) # This is syntactic-sugar, so that when calling this decorator without args, you can either diff --git a/flytekit/core/base_task.py b/flytekit/core/base_task.py index 570e1c4d0c..53fa011185 100644 --- a/flytekit/core/base_task.py +++ b/flytekit/core/base_task.py @@ -72,6 +72,7 @@ class TaskMetadata(object): Args: cache (bool): Indicates if caching should be enabled. See :std:ref:`Caching ` + cache_serialize (bool): Indicates if identical (ie. same inputs) instances of this task should be executed in serial when caching is enabled. See :std:ref:`Caching ` cache_version (str): Version to be used for the cached value interruptible (Optional[bool]): Indicates that this task can be interrupted and/or scheduled on nodes with lower QoS guarantees that can include pre-emption. This can reduce the monetary cost executions incur at the @@ -85,6 +86,7 @@ class TaskMetadata(object): """ cache: bool = False + cache_serialize: bool = False cache_version: str = "" interruptible: Optional[bool] = None deprecated: str = "" @@ -99,6 +101,8 @@ def __post_init__(self): raise ValueError("timeout should be duration represented as either a datetime.timedelta or int seconds") if self.cache and not self.cache_version: raise ValueError("Caching is enabled ``cache=True`` but ``cache_version`` is not set.") + if self.cache_serialize and not self.cache: + raise ValueError("Cache serialize is enabled ``cache_serialize=True`` but ``cache`` is not enabled.") @property def retry_strategy(self) -> _literal_models.RetryStrategy: @@ -120,6 +124,7 @@ def to_taskmetadata_model(self) -> _task_model.TaskMetadata: interruptible=self.interruptible, discovery_version=self.cache_version, deprecated_error_message=self.deprecated, + cache_serializable=self.cache_serialize, ) diff --git a/flytekit/core/task.py b/flytekit/core/task.py index 45b7e2dc94..e070359c98 100644 --- a/flytekit/core/task.py +++ b/flytekit/core/task.py @@ -75,6 +75,7 @@ def task( _task_function: Optional[Callable] = None, task_config: Optional[Any] = None, cache: bool = False, + cache_serialize: bool = False, cache_version: str = "", retries: int = 0, interruptible: Optional[bool] = None, @@ -121,6 +122,10 @@ def my_task(x: int, y: typing.Dict[str, str]) -> str: :param task_config: This argument provides configuration for a specific task types. Please refer to the plugins documentation for the right object to use. :param cache: Boolean that indicates if caching should be enabled + :param cache_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. This + parameter does nothing without also setting the cache parameter. :param cache_version: Cache version to use. Changes to the task signature will automatically trigger a cache miss, but you can always manually update this field as well to force a cache miss. You should also manually bump this version if the function body/business logic has changed, but the signature hasn't. @@ -176,6 +181,7 @@ def foo2(): def wrapper(fn) -> PythonFunctionTask: _metadata = TaskMetadata( cache=cache, + cache_serialize=cache_serialize, cache_version=cache_version, retries=retries, interruptible=interruptible, diff --git a/flytekit/models/task.py b/flytekit/models/task.py index c20872096d..76caa0444d 100644 --- a/flytekit/models/task.py +++ b/flytekit/models/task.py @@ -180,6 +180,7 @@ def __init__( interruptible, discovery_version, deprecated_error_message, + cache_serializable, ): """ Information needed at runtime to determine behavior such as whether or not outputs are discoverable, timeouts, @@ -197,6 +198,8 @@ def __init__( task are the same and the discovery_version is also the same. :param Text deprecated: This string can be used to mark the task as deprecated. Consumers of the task will receive deprecation warnings. + :param bool cache_serializable: Whether or not caching operations are executed in serial. This means only a + single instance over identical inputs is executed, other concurrent executions wait for the cached results. """ self._discoverable = discoverable self._runtime = runtime @@ -205,6 +208,7 @@ def __init__( self._retries = retries self._discovery_version = discovery_version self._deprecated_error_message = deprecated_error_message + self._cache_serializable = cache_serializable @property def discoverable(self): @@ -265,6 +269,15 @@ def deprecated_error_message(self): """ return self._deprecated_error_message + @property + def cache_serializable(self): + """ + Whether or not caching operations are executed in serial. This means only a single instance over identical + inputs is executed, other concurrent executions wait for the cached results. + :rtype: bool + """ + return self._cache_serializable + def to_flyte_idl(self): """ :rtype: flyteidl.admin.task_pb2.TaskMetadata @@ -276,6 +289,7 @@ def to_flyte_idl(self): interruptible=self.interruptible, discovery_version=self.discovery_version, deprecated_error_message=self.deprecated_error_message, + cache_serializable=self.cache_serializable, ) if self.timeout: tm.timeout.FromTimedelta(self.timeout) @@ -295,6 +309,7 @@ def from_flyte_idl(cls, pb2_object): retries=_literals.RetryStrategy.from_flyte_idl(pb2_object.retries), discovery_version=pb2_object.discovery_version, deprecated_error_message=pb2_object.deprecated_error_message, + cache_serializable=pb2_object.cache_serializable, ) diff --git a/flytekit/sdk/tasks.py b/flytekit/sdk/tasks.py index dd43055077..0c73ad66b7 100644 --- a/flytekit/sdk/tasks.py +++ b/flytekit/sdk/tasks.py @@ -137,6 +137,7 @@ def python_task( cache=False, timeout=None, environment=None, + cache_serialize=False, cls=None, ): """ @@ -224,6 +225,10 @@ def my_task(wf_params, int_list, sum_of_list): :param dict[Text,Text] environment: [optional] environment variables to set when executing this task. + :param bool cache_serialize: [optional] boolean describing if instances of this cachable task should be executed + in serial. This means only a single instances executes and other concurrent executions wait for it to complete + and reuse the cached outputs. + :param cls: This can be used to override the task implementation with a user-defined extension. The class provided must be a subclass of flytekit.common.tasks.sdk_runnable.SdkRunnableTask. A user can use this to inject bespoke logic into the base Flyte programming model. @@ -250,6 +255,7 @@ def wrapper(fn): discoverable=cache, timeout=timeout or _datetime.timedelta(seconds=0), environment=environment, + cache_serializable=cache_serialize, custom={}, ) @@ -278,6 +284,7 @@ def dynamic_task( allowed_failure_ratio=None, max_concurrency=None, environment=None, + cache_serialize=False, cls=None, ): """ @@ -375,6 +382,9 @@ def my_task(wf_params, out): This is a stand-in pending better concurrency controls for special use-cases. The existence of this parameter is not guaranteed between versions and therefore it is NOT recommended that it be used. :param dict[Text,Text] environment: [optional] environment variables to set when executing this task. + :param bool cache_serialize: [optional] boolean describing if instances of this cachable task should be executed + in serial. This means only a single instances executes and other concurrent executions wait for it to complete + and reuse the cached outputs. :param cls: This can be used to override the task implementation with a user-defined extension. The class provided must be a subclass of flytekit.common.tasks.sdk_runnable.SdkRunnableTask. Generally, it should be a subclass of flytekit.common.tasks.sdk_dynamic.SdkDynamicTask. A user can use this parameter to inject bespoke @@ -403,6 +413,7 @@ def wrapper(fn): allowed_failure_ratio=allowed_failure_ratio, max_concurrency=max_concurrency, environment=environment or {}, + cache_serializable=cache_serialize, custom={}, ) @@ -423,6 +434,7 @@ def spark_task( spark_conf=None, hadoop_conf=None, environment=None, + cache_serialize=False, cls=None, ): """ @@ -471,6 +483,9 @@ def sparky(wf_params, spark_context, a): :param dict[Text,Text] spark_conf: A definition of key-value pairs for spark config for the job. :param dict[Text,Text] hadoop_conf: A definition of key-value pairs for hadoop config for the job. :param dict[Text,Text] environment: [optional] environment variables to set when executing this task. + :param bool cache_serialize: [optional] boolean describing if instances of this cachable task should be executed + in serial. This means only a single instances executes and other concurrent executions wait for it to complete + and reuse the cached outputs. :param cls: This can be used to override the task implementation with a user-defined extension. The class provided must be a subclass of flytekit.common.tasks.sdk_runnable.SdkRunnableTask. Generally, it should be a subclass of flytekit.common.tasks.spark_task.SdkSparkTask. A user can use this parameter to inject bespoke @@ -492,6 +507,7 @@ def wrapper(fn): spark_conf=spark_conf or {}, hadoop_conf=hadoop_conf or {}, environment=environment or {}, + cache_serializable=cache_serialize, ) if _task_function: @@ -514,6 +530,7 @@ def generic_spark_task( spark_conf=None, hadoop_conf=None, environment=None, + cache_serialize=False, ): """ Create a generic spark task. This task will connect to a Spark cluster, configure the environment, @@ -536,6 +553,7 @@ def generic_spark_task( spark_conf=spark_conf or {}, hadoop_conf=hadoop_conf or {}, environment=environment or {}, + cache_serializable=cache_serialize, ) @@ -563,6 +581,7 @@ def hive_task( cache=False, timeout=None, environment=None, + cache_serialize=False, cls=None, ): """ @@ -648,6 +667,9 @@ def test_hive(wf_params, a): indefinitely. If a null timedelta is passed (i.e. timedelta(seconds=0)), the task will not timeout. :param dict[Text,Text] environment: Environment variables to set for the execution of the query-generating container. + :param bool cache_serialize: [optional] boolean describing if instances of this cachable task should be executed + in serial. This means only a single instances executes and other concurrent executions wait for it to complete + and reuse the cached outputs. :param cls: This can be used to override the task implementation with a user-defined extension. The class provided should be a subclass of flytekit.common.tasks.sdk_runnable.SdkRunnableTask. Generally, it should be a subclass of flytekit.common.tasks.hive_task.SdkHiveTask. A user can use this to inject bespoke logic into @@ -678,6 +700,7 @@ def wrapper(fn): cluster_label="", tags=[], environment=environment or {}, + cache_serializable=cache_serialize, ) if _task_function: @@ -705,6 +728,7 @@ def qubole_hive_task( cluster_label=None, tags=None, environment=None, + cache_serialize=False, cls=None, ): """ @@ -793,6 +817,9 @@ def test_hive(wf_params, a): passed to Qubole. :param dict[Text,Text] environment: Environment variables to set for the execution of the query-generating container. + :param bool cache_serialize: [optional] boolean describing if instances of this cachable task should be executed + in serial. This means only a single instances executes and other concurrent executions wait for it to complete + and reuse the cached outputs. :param cls: This can be used to override the task implementation with a user-defined extension. The class provided should be a subclass of flytekit.common.tasks.sdk_runnable.SdkRunnableTask. Generally, it should be a subclass of flytekit.common.tasks.hive_task.SdkHiveTask. A user can use this to inject bespoke logic into @@ -823,6 +850,7 @@ def wrapper(fn): cluster_label=cluster_label or "", tags=tags or [], environment=environment or {}, + cache_serializable=cache_serialize, ) # This is syntactic-sugar, so that when calling this decorator without args, you can either @@ -850,6 +878,7 @@ def sidecar_task( cache=False, timeout=None, environment=None, + cache_serialize=False, pod_spec=None, primary_container_name=None, annotations=None, @@ -977,6 +1006,10 @@ def a_sidecar_task(wfparams): :param dict[Text,Text] environment: [optional] environment variables to set when executing this task. + :param bool cache_serialize: [optional] boolean describing if instances of this cachable task should be executed + in serial. This means only a single instances executes and other concurrent executions wait for it to complete + and reuse the cached outputs. + :param k8s.io.api.core.v1.generated_pb2.PodSpec pod_spec: [optional] PodSpec to bring up alongside task execution. :param Text primary_container_name: primary container to monitor for the duration of the task. @@ -1013,6 +1046,7 @@ def wrapper(fn): discoverable=cache, timeout=timeout or _datetime.timedelta(seconds=0), environment=environment, + cache_serializable=cache_serialize, pod_spec=pod_spec, primary_container_name=primary_container_name, annotations=annotations, @@ -1044,6 +1078,7 @@ def dynamic_sidecar_task( allowed_failure_ratio=None, max_concurrency=None, environment=None, + cache_serialize=False, pod_spec=None, primary_container_name=None, annotations=None, @@ -1161,6 +1196,9 @@ def my_task(wf_params, out): This is a stand-in pending better concurrency controls for special use-cases. The existence of this parameter is not guaranteed between versions and therefore it is NOT recommended that it be used. :param dict[Text,Text] environment: [optional] environment variables to set when executing this task. + :param bool cache_serialize: [optional] boolean describing if instances of this cachable task should be executed + in serial. This means only a single instances executes and other concurrent executions wait for it to complete + and reuse the cached outputs. :param k8s.io.api.core.v1.generated_pb2.PodSpec pod_spec: PodSpec to bring up alongside task execution. :param Text primary_container_name: primary container to monitor for the duration of the task. :param dict[Text, Text] annotations: [optional] kubernetes annotations @@ -1193,6 +1231,7 @@ def wrapper(fn): allowed_failure_ratio=allowed_failure_ratio, max_concurrency=max_concurrency, environment=environment, + cache_serializable=cache_serialize, pod_spec=pod_spec, primary_container_name=primary_container_name, annotations=annotations, diff --git a/requirements-spark2.txt b/requirements-spark2.txt index 37a2664140..a0457279bc 100644 --- a/requirements-spark2.txt +++ b/requirements-spark2.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.8 +# This file is autogenerated by pip-compile with python 3.9 # To update, run: # # make requirements-spark2.txt @@ -22,13 +22,13 @@ bcrypt==3.2.0 # via paramiko binaryornot==0.4.4 # via cookiecutter -black==21.10b0 +black==21.11b1 # via papermill bleach==4.1.0 # via nbconvert -boto3==1.20.2 +boto3==1.20.17 # via sagemaker-training -botocore==1.23.2 +botocore==1.23.17 # via # boto3 # s3transfer @@ -41,7 +41,7 @@ cffi==1.15.0 # pynacl chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.7 +charset-normalizer==2.0.8 # via requests checksumdir==1.2.0 # via flytekit @@ -58,7 +58,7 @@ cookiecutter==1.7.3 # via flytekit croniter==1.0.15 # via flytekit -cryptography==35.0.0 +cryptography==36.0.0 # via # paramiko # secretstorage @@ -72,24 +72,24 @@ defusedxml==0.7.1 # via nbconvert deprecated==1.2.13 # via flytekit -diskcache==5.2.1 +diskcache==5.3.0 # via flytekit docker-image-py==0.1.12 # via flytekit -docstring-parser==0.12 +docstring-parser==0.13 # via flytekit entrypoints==0.3 # via # jupyter-client # nbconvert # papermill -flyteidl==0.21.8 +flyteidl==0.21.11 # via flytekit gevent==21.8.0 # via sagemaker-training greenlet==1.1.2 # via gevent -grpcio==1.41.1 +grpcio==1.42.0 # via flytekit hmsclient==0.1.1 # via flytekit @@ -101,13 +101,13 @@ inotify_simple==1.2.1 # via sagemaker-training ipykernel==5.5.6 # via flytekit -ipython==7.29.0 +ipython==7.30.0 # via ipykernel ipython-genutils==0.2.0 # via # ipykernel # nbformat -jedi==0.18.0 +jedi==0.18.1 # via ipython jeepney==0.7.1 # via @@ -128,7 +128,7 @@ jsonschema==3.2.0 # via # -r requirements.in # nbformat -jupyter-client==7.0.6 +jupyter-client==7.1.0 # via # ipykernel # nbclient @@ -141,11 +141,11 @@ jupyterlab-pygments==0.1.2 # via nbconvert k8s-proto==0.0.3 # via flytekit -keyring==23.2.1 +keyring==23.4.0 # via flytekit markupsafe==2.0.1 # via jinja2 -marshmallow==3.14.0 +marshmallow==3.14.1 # via # dataclasses-json # marshmallow-enum @@ -164,18 +164,18 @@ mypy-extensions==0.4.3 # typing-inspect natsort==8.0.0 # via flytekit -nbclient==0.5.5 +nbclient==0.5.9 # via # nbconvert # papermill -nbconvert==6.2.0 +nbconvert==6.3.0 # via flytekit nbformat==5.1.3 # via # nbclient # nbconvert # papermill -nest-asyncio==1.5.1 +nest-asyncio==1.5.2 # via # jupyter-client # nbclient @@ -186,7 +186,7 @@ numpy==1.21.4 # pyarrow # sagemaker-training # scipy -packaging==21.2 +packaging==21.3 # via bleach pandas==1.3.4 # via flytekit @@ -194,9 +194,9 @@ pandocfilters==1.5.0 # via nbconvert papermill==2.3.3 # via flytekit -paramiko==2.8.0 +paramiko==2.8.1 # via sagemaker-training -parso==0.8.2 +parso==0.8.3 # via jedi pathspec==0.9.0 # via black @@ -208,7 +208,7 @@ platformdirs==2.4.0 # via black poyo==0.5.0 # via cookiecutter -prompt-toolkit==3.0.22 +prompt-toolkit==3.0.23 # via ipython protobuf==3.19.1 # via @@ -224,7 +224,7 @@ py==1.11.0 # via retry py4j==0.10.9.2 # via pyspark -pyarrow==6.0.0 +pyarrow==6.0.1 # via flytekit pycparser==2.21 # via cffi @@ -235,7 +235,7 @@ pygments==2.10.0 # nbconvert pynacl==1.4.0 # via paramiko -pyparsing==2.4.7 +pyparsing==3.0.6 # via packaging pyrsistent==0.18.0 # via jsonschema @@ -265,7 +265,7 @@ pyyaml==5.4.1 # papermill pyzmq==22.3.0 # via jupyter-client -regex==2021.11.9 +regex==2021.11.10 # via # black # docker-image-py @@ -275,7 +275,7 @@ requests==2.26.0 # flytekit # papermill # responses -responses==0.15.0 +responses==0.16.0 # via flytekit retry==0.9.2 # via flytekit @@ -285,7 +285,7 @@ s3transfer==0.5.0 # via boto3 sagemaker-training==3.9.2 # via flytekit -scipy==1.7.2 +scipy==1.7.3 # via sagemaker-training secretstorage==3.3.1 # via keyring @@ -335,7 +335,7 @@ traitlets==5.1.1 # nbclient # nbconvert # nbformat -typing-extensions==3.10.0.2 +typing-extensions==4.0.1 # via # black # typing-inspect diff --git a/requirements.txt b/requirements.txt index ee632aa7d5..344af69890 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.8 +# This file is autogenerated by pip-compile with python 3.9 # To update, run: # # make requirements.txt @@ -20,13 +20,13 @@ bcrypt==3.2.0 # via paramiko binaryornot==0.4.4 # via cookiecutter -black==21.10b0 +black==21.11b1 # via papermill bleach==4.1.0 # via nbconvert -boto3==1.20.2 +boto3==1.20.17 # via sagemaker-training -botocore==1.23.2 +botocore==1.23.17 # via # boto3 # s3transfer @@ -39,7 +39,7 @@ cffi==1.15.0 # pynacl chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.7 +charset-normalizer==2.0.8 # via requests checksumdir==1.2.0 # via flytekit @@ -56,7 +56,7 @@ cookiecutter==1.7.3 # via flytekit croniter==1.0.15 # via flytekit -cryptography==35.0.0 +cryptography==36.0.0 # via # paramiko # secretstorage @@ -70,24 +70,24 @@ defusedxml==0.7.1 # via nbconvert deprecated==1.2.13 # via flytekit -diskcache==5.2.1 +diskcache==5.3.0 # via flytekit docker-image-py==0.1.12 # via flytekit -docstring-parser==0.12 +docstring-parser==0.13 # via flytekit entrypoints==0.3 # via # jupyter-client # nbconvert # papermill -flyteidl==0.21.8 +flyteidl==0.21.11 # via flytekit gevent==21.8.0 # via sagemaker-training greenlet==1.1.2 # via gevent -grpcio==1.41.1 +grpcio==1.42.0 # via flytekit hmsclient==0.1.1 # via flytekit @@ -99,13 +99,13 @@ inotify_simple==1.2.1 # via sagemaker-training ipykernel==5.5.6 # via flytekit -ipython==7.29.0 +ipython==7.30.0 # via ipykernel ipython-genutils==0.2.0 # via # ipykernel # nbformat -jedi==0.18.0 +jedi==0.18.1 # via ipython jeepney==0.7.1 # via @@ -126,7 +126,7 @@ jsonschema==3.2.0 # via # -r requirements.in # nbformat -jupyter-client==7.0.6 +jupyter-client==7.1.0 # via # ipykernel # nbclient @@ -139,11 +139,11 @@ jupyterlab-pygments==0.1.2 # via nbconvert k8s-proto==0.0.3 # via flytekit -keyring==23.2.1 +keyring==23.4.0 # via flytekit markupsafe==2.0.1 # via jinja2 -marshmallow==3.14.0 +marshmallow==3.14.1 # via # dataclasses-json # marshmallow-enum @@ -162,18 +162,18 @@ mypy-extensions==0.4.3 # typing-inspect natsort==8.0.0 # via flytekit -nbclient==0.5.5 +nbclient==0.5.9 # via # nbconvert # papermill -nbconvert==6.2.0 +nbconvert==6.3.0 # via flytekit nbformat==5.1.3 # via # nbclient # nbconvert # papermill -nest-asyncio==1.5.1 +nest-asyncio==1.5.2 # via # jupyter-client # nbclient @@ -184,7 +184,7 @@ numpy==1.21.4 # pyarrow # sagemaker-training # scipy -packaging==21.2 +packaging==21.3 # via bleach pandas==1.3.4 # via flytekit @@ -192,9 +192,9 @@ pandocfilters==1.5.0 # via nbconvert papermill==2.3.3 # via flytekit -paramiko==2.8.0 +paramiko==2.8.1 # via sagemaker-training -parso==0.8.2 +parso==0.8.3 # via jedi pathspec==0.9.0 # via black @@ -206,7 +206,7 @@ platformdirs==2.4.0 # via black poyo==0.5.0 # via cookiecutter -prompt-toolkit==3.0.22 +prompt-toolkit==3.0.23 # via ipython protobuf==3.19.1 # via @@ -222,7 +222,7 @@ py==1.11.0 # via retry py4j==0.10.9.2 # via pyspark -pyarrow==6.0.0 +pyarrow==6.0.1 # via flytekit pycparser==2.21 # via cffi @@ -233,7 +233,7 @@ pygments==2.10.0 # nbconvert pynacl==1.4.0 # via paramiko -pyparsing==2.4.7 +pyparsing==3.0.6 # via packaging pyrsistent==0.18.0 # via jsonschema @@ -263,7 +263,7 @@ pyyaml==5.4.1 # papermill pyzmq==22.3.0 # via jupyter-client -regex==2021.11.9 +regex==2021.11.10 # via # black # docker-image-py @@ -273,7 +273,7 @@ requests==2.26.0 # flytekit # papermill # responses -responses==0.15.0 +responses==0.16.0 # via flytekit retry==0.9.2 # via flytekit @@ -283,7 +283,7 @@ s3transfer==0.5.0 # via boto3 sagemaker-training==3.9.2 # via flytekit -scipy==1.7.2 +scipy==1.7.3 # via sagemaker-training secretstorage==3.3.1 # via keyring @@ -333,7 +333,7 @@ traitlets==5.1.1 # nbclient # nbconvert # nbformat -typing-extensions==3.10.0.2 +typing-extensions==4.0.1 # via # black # typing-inspect diff --git a/tests/flytekit/common/parameterizers.py b/tests/flytekit/common/parameterizers.py index c7d022d860..33bd7712ee 100644 --- a/tests/flytekit/common/parameterizers.py +++ b/tests/flytekit/common/parameterizers.py @@ -124,8 +124,9 @@ interruptible, discovery_version, deprecated, + cache_serializable, ) - for discoverable, runtime_metadata, timeout, retry_strategy, interruptible, discovery_version, deprecated in product( + for discoverable, runtime_metadata, timeout, retry_strategy, interruptible, discovery_version, deprecated, cache_serializable in product( [True, False], LIST_OF_RUNTIME_METADATA, [timedelta(days=i) for i in range(3)], @@ -133,6 +134,7 @@ LIST_OF_INTERRUPTIBLE, ["1.0"], ["deprecated"], + [True, False], ) ] diff --git a/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt b/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt index a37b0c3af8..f9ecc1e38d 100644 --- a/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt +++ b/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.8 +# This file is autogenerated by pip-compile with python 3.9 # To update, run: # # make tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt @@ -14,7 +14,7 @@ cffi==1.15.0 # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.7 +charset-normalizer==2.0.8 # via requests checksumdir==1.2.0 # via flytekit @@ -28,7 +28,7 @@ cookiecutter==1.7.3 # via flytekit croniter==1.0.15 # via flytekit -cryptography==35.0.0 +cryptography==36.0.0 # via secretstorage cycler==0.11.0 # via matplotlib @@ -38,17 +38,19 @@ decorator==5.1.0 # via retry deprecated==1.2.13 # via flytekit -diskcache==5.2.1 +diskcache==5.3.0 # via flytekit docker-image-py==0.1.12 # via flytekit -docstring-parser==0.12 +docstring-parser==0.13 # via flytekit -flyteidl==0.21.8 +flyteidl==0.21.11 # via flytekit flytekit==0.24.0 # via -r tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.in -grpcio==1.41.1 +fonttools==4.28.2 + # via matplotlib +grpcio==1.42.0 # via flytekit idna==3.3 # via requests @@ -66,13 +68,13 @@ jinja2-time==0.2.0 # via cookiecutter joblib==1.1.0 # via -r tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.in -keyring==23.2.1 +keyring==23.4.0 # via flytekit kiwisolver==1.3.2 # via matplotlib markupsafe==2.0.1 # via jinja2 -marshmallow==3.14.0 +marshmallow==3.14.1 # via # dataclasses-json # marshmallow-enum @@ -81,7 +83,7 @@ marshmallow-enum==1.5.1 # via dataclasses-json marshmallow-jsonschema==0.13.0 # via flytekit -matplotlib==3.4.3 +matplotlib==3.5.0 # via -r tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.in mypy-extensions==0.4.3 # via typing-inspect @@ -93,8 +95,12 @@ numpy==1.21.4 # opencv-python # pandas # pyarrow -opencv-python==4.5.4.58 +opencv-python==4.5.4.60 # via -r tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.in +packaging==21.3 + # via + # matplotlib + # setuptools-scm pandas==1.3.4 # via flytekit pillow==8.4.0 @@ -107,12 +113,14 @@ protobuf==3.19.1 # flytekit py==1.11.0 # via retry -pyarrow==6.0.0 +pyarrow==6.0.1 # via flytekit pycparser==2.21 # via cffi -pyparsing==3.0.5 - # via matplotlib +pyparsing==3.0.6 + # via + # matplotlib + # packaging python-dateutil==2.8.1 # via # arrow @@ -130,19 +138,21 @@ pytz==2018.4 # via # flytekit # pandas -regex==2021.11.9 +regex==2021.11.10 # via docker-image-py requests==2.26.0 # via # cookiecutter # flytekit # responses -responses==0.15.0 +responses==0.16.0 # via flytekit retry==0.9.2 # via flytekit secretstorage==3.3.1 # via keyring +setuptools-scm==6.3.2 + # via matplotlib six==1.16.0 # via # cookiecutter @@ -156,7 +166,9 @@ statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -typing-extensions==3.10.0.2 +tomli==1.2.2 + # via setuptools-scm +typing-extensions==4.0.1 # via typing-inspect typing-inspect==0.7.1 # via dataclasses-json @@ -175,3 +187,6 @@ wrapt==1.13.3 # flytekit zipp==3.6.0 # via importlib-metadata + +# The following packages are considered to be unsafe in a requirements file: +# setuptools diff --git a/tests/flytekit/unit/common_tests/tasks/test_sdk_runnable.py b/tests/flytekit/unit/common_tests/tasks/test_sdk_runnable.py index 80d5b20421..43e677551c 100644 --- a/tests/flytekit/unit/common_tests/tasks/test_sdk_runnable.py +++ b/tests/flytekit/unit/common_tests/tasks/test_sdk_runnable.py @@ -29,6 +29,7 @@ def add_one(wf_params, value_in, value_out): False, None, {}, + False, None, ) t.add_inputs({"value_in": interface.Variable(primitives.Integer.to_flyte_literal_type(), "")}) diff --git a/tests/flytekit/unit/common_tests/test_workflow_promote.py b/tests/flytekit/unit/common_tests/test_workflow_promote.py index cdd1231107..f165f4c231 100644 --- a/tests/flytekit/unit/common_tests/test_workflow_promote.py +++ b/tests/flytekit/unit/common_tests/test_workflow_promote.py @@ -60,6 +60,7 @@ def get_sample_task_metadata(): True, "0.1.1b0", "This is deprecated!", + True, ) diff --git a/tests/flytekit/unit/core/test_python_function_task.py b/tests/flytekit/unit/core/test_python_function_task.py index 7101b585d3..708bacfe50 100644 --- a/tests/flytekit/unit/core/test_python_function_task.py +++ b/tests/flytekit/unit/core/test_python_function_task.py @@ -89,10 +89,36 @@ def foo(i: int): def test_metadata(): - @task(cache=True, cache_version="1.0") + # test cache, cache_serialize, and cache_version are correctly set + @task(cache=True, cache_serialize=True, cache_version="1.0") def foo(i: str): print(f"{i}") - metadata = foo.metadata - assert metadata.cache is True - assert metadata.cache_version == "1.0" + foo_metadata = foo.metadata + assert foo_metadata.cache is True + assert foo_metadata.cache_serialize is True + assert foo_metadata.cache_version == "1.0" + + # test cache, cache_serialize, and cache_version at no unecessarily set + @task() + def bar(i: str): + print(f"{i}") + + bar_metadata = bar.metadata + assert bar_metadata.cache is False + assert bar_metadata.cache_serialize is False + assert bar_metadata.cache_version == "" + + # test missing cache_version + with pytest.raises(ValueError): + + @task(cache=True) + def foo_missing_cache_version(i: str): + print(f"{i}") + + # test missing cache + with pytest.raises(ValueError): + + @task(cache_serialize=True) + def foo_missing_cache(i: str): + print(f"{i}") diff --git a/tests/flytekit/unit/models/test_tasks.py b/tests/flytekit/unit/models/test_tasks.py index f3ae95fb3a..dec28f0f54 100644 --- a/tests/flytekit/unit/models/test_tasks.py +++ b/tests/flytekit/unit/models/test_tasks.py @@ -70,6 +70,7 @@ def test_task_metadata(): True, "0.1.1b0", "This is deprecated!", + True, ) assert obj.discoverable is True @@ -136,6 +137,7 @@ def test_task_template__k8s_pod_target(): False, "1.0", "deprecated", + False, ), interface_models.TypedInterface( # inputs diff --git a/tests/flytekit/unit/models/test_workflow_closure.py b/tests/flytekit/unit/models/test_workflow_closure.py index 3e19a80657..3a42f5af81 100644 --- a/tests/flytekit/unit/models/test_workflow_closure.py +++ b/tests/flytekit/unit/models/test_workflow_closure.py @@ -35,6 +35,7 @@ def test_workflow_closure(): True, "0.1.1b0", "This is deprecated!", + True, ) cpu_resource = _task.Resources.ResourceEntry(_task.Resources.ResourceName.CPU, "1") From 687851a99070f03cab13d99db59318aeedde07e7 Mon Sep 17 00:00:00 2001 From: bstadlbauer <11799671+bstadlbauer@users.noreply.github.com> Date: Fri, 17 Dec 2021 20:29:26 +0100 Subject: [PATCH 040/128] =?UTF-8?q?When=20using=20the=20`task`=20and=20`wo?= =?UTF-8?q?rkflow`=20decorator,=20correctly=20wrap=20the=20fu=E2=80=A6=20(?= =?UTF-8?q?#780)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * When using the `task` and `workflow` decorator, correctly wrap the function This enables tooling such as docstring search tools to unwrap the object and show the correct docstring. Signed-off-by: Bernhard Stadlbauer * Remove blackshark copyright header Signed-off-by: Bernhard Stadlbauer * Fix broken great expectations test Signed-off-by: Bernhard Stadlbauer * Add test for stacked decorators Signed-off-by: Bernhard Stadlbauer Co-authored-by: Bernhard Stadlbauer Co-authored-by: Bernhard Stadlbauer Signed-off-by: maximsmol --- flytekit/core/task.py | 3 +- flytekit/core/workflow.py | 2 + .../tests/test_task.py | 4 +- tests/flytekit/unit/core/test_wrapping.py | 58 +++++++++++++++++++ 4 files changed, 64 insertions(+), 3 deletions(-) create mode 100644 tests/flytekit/unit/core/test_wrapping.py diff --git a/flytekit/core/task.py b/flytekit/core/task.py index e070359c98..a5f904bf07 100644 --- a/flytekit/core/task.py +++ b/flytekit/core/task.py @@ -1,4 +1,5 @@ import datetime as _datetime +from functools import update_wrapper from typing import Any, Callable, Dict, List, Optional, Type, Union from flytekit.core.base_task import TaskMetadata, TaskResolverMixin @@ -201,7 +202,7 @@ def wrapper(fn) -> PythonFunctionTask: execution_mode=execution_mode, task_resolver=task_resolver, ) - + update_wrapper(task_instance, fn) return task_instance if _task_function: diff --git a/flytekit/core/workflow.py b/flytekit/core/workflow.py index 744ecfbb11..ffa6aae934 100644 --- a/flytekit/core/workflow.py +++ b/flytekit/core/workflow.py @@ -2,6 +2,7 @@ from dataclasses import dataclass from enum import Enum +from functools import update_wrapper from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union from flytekit.common import constants as _common_constants @@ -730,6 +731,7 @@ def wrapper(fn): docstring=Docstring(callable_=fn), ) workflow_instance.compile() + update_wrapper(workflow_instance, fn) return workflow_instance if _workflow_function: diff --git a/plugins/flytekit-greatexpectations/tests/test_task.py b/plugins/flytekit-greatexpectations/tests/test_task.py index 9906085db1..578e9ee791 100644 --- a/plugins/flytekit-greatexpectations/tests/test_task.py +++ b/plugins/flytekit-greatexpectations/tests/test_task.py @@ -149,7 +149,6 @@ def valid_wf(dataset: str = "yellow_tripdata_sample_2019-01.csv") -> int: task_object(dataset=dataset) return my_task(csv_file=dataset) - @pytest.mark.xfail(strict=True) @workflow def invalid_wf(dataset: str = "yellow_tripdata_sample_2019-02.csv") -> int: task_object(dataset=dataset) @@ -158,7 +157,8 @@ def invalid_wf(dataset: str = "yellow_tripdata_sample_2019-02.csv") -> int: valid_result = valid_wf() assert valid_result == 10000 - invalid_wf() + with pytest.raises(ValidationError, match=r".*passenger_count -> expect_column_min_to_be_between.*"): + invalid_wf() def test_ge_workflow(): diff --git a/tests/flytekit/unit/core/test_wrapping.py b/tests/flytekit/unit/core/test_wrapping.py new file mode 100644 index 0000000000..97ff2bafec --- /dev/null +++ b/tests/flytekit/unit/core/test_wrapping.py @@ -0,0 +1,58 @@ +from functools import wraps + +from flytekit import task, workflow + + +def test_task_correctly_wrapped(): + @task + def my_task(a: int) -> int: + return a + + assert my_task.__wrapped__ == my_task._task_function + + +def test_stacked_decorators(): + def task_decorator_1(fn): + @wraps(fn) + def wrapper(*args, **kwargs): + print("running task_decorator_1") + return fn(*args, **kwargs) + + return wrapper + + def task_decorator_2(fn): + @wraps(fn) + def wrapper(*args, **kwargs): + print("running task_decorator_2") + return fn(*args, **kwargs) + + return wrapper + + def task_decorator_3(fn): + @wraps(fn) + def wrapper(*args, **kwargs): + print("running task_decorator_3") + return fn(*args, **kwargs) + + return wrapper + + @task + @task_decorator_1 + @task_decorator_2 + @task_decorator_3 + def my_task(x: int) -> int: + """Some function doc""" + print("running my_task") + return x + 1 + + assert my_task.__wrapped__.__doc__ == "Some function doc" + assert my_task.__wrapped__ == my_task._task_function + assert my_task(x=10) == 11 + + +def test_wf_correctly_wrapped(): + @workflow + def my_workflow(a: int) -> int: + return a + + assert my_workflow.__wrapped__ == my_workflow._workflow_function From 669593e4c52803c3a6e870b849560e6187ab3aca Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Mon, 20 Dec 2021 12:23:35 -0800 Subject: [PATCH 041/128] Add option to flyte-cli for specifying root certificate (#783) Signed-off-by: Yee Hing Tong Signed-off-by: maximsmol --- flytekit/clients/raw.py | 14 ++- flytekit/clis/flyte_cli/main.py | 146 ++++++++++++++-------- tests/flytekit/unit/remote/test_remote.py | 2 +- 3 files changed, 105 insertions(+), 57 deletions(-) diff --git a/flytekit/clients/raw.py b/flytekit/clients/raw.py index 25c3b75b9b..b8f7f578cf 100644 --- a/flytekit/clients/raw.py +++ b/flytekit/clients/raw.py @@ -183,7 +183,7 @@ class RawSynchronousFlyteClient(object): be explicit as opposed to inferred from the environment or a configuration file. """ - def __init__(self, url, insecure=False, credentials=None, options=None): + def __init__(self, url, insecure=False, credentials=None, options=None, root_cert_file=None): """ Initializes a gRPC channel to the given Flyte Admin service. @@ -192,8 +192,7 @@ def __init__(self, url, insecure=False, credentials=None, options=None): :param Text credentials: [Optional] If provided, a secure channel will be opened with the Flyte Admin Service. :param dict[Text, Text] options: [Optional] A dict of key-value string pairs for configuring the gRPC core runtime. - :param list[(Text,Text)] metadata: [Optional] metadata pairs to be transmitted to the - service-side of the RPC. + :param root_cert_file: Path to a local certificate file if you want. """ self._channel = None self._url = url @@ -201,9 +200,16 @@ def __init__(self, url, insecure=False, credentials=None, options=None): if insecure: self._channel = _insecure_channel(url, options=list((options or {}).items())) else: + if root_cert_file: + with open(root_cert_file, "rb") as fh: + cert_bytes = fh.read() + channel_creds = _ssl_channel_credentials(root_certificates=cert_bytes) + else: + channel_creds = _ssl_channel_credentials() + self._channel = _secure_channel( url, - credentials or _ssl_channel_credentials(), + credentials or channel_creds, options=list((options or {}).items()), ) self._stub = _admin_service.AdminServiceStub(self._channel) diff --git a/flytekit/clis/flyte_cli/main.py b/flytekit/clis/flyte_cli/main.py index 4e0390a041..4dab451827 100644 --- a/flytekit/clis/flyte_cli/main.py +++ b/flytekit/clis/flyte_cli/main.py @@ -8,7 +8,6 @@ import click as _click import requests as _requests -import six as _six from flyteidl.admin import launch_plan_pb2 as _launch_plan_pb2 from flyteidl.admin import task_pb2 as _task_pb2 from flyteidl.admin import workflow_pb2 as _workflow_pb2 @@ -66,7 +65,7 @@ except ImportError: # Python 2 import urlparse as _urlparse -_tt = _six.text_type +_tt = str # Similar to how kubectl has a config file in the users home directory, this Flyte CLI will also look for one. # The format of this config file is the same as a workflow's config file, except that the relevant fields are different. @@ -129,7 +128,7 @@ def _get_io_string(literal_map, verbose=False): v.verbose_string() if verbose else v.short_string(), ), ) - for k, v in _six.iteritems(value_dict) + for k, v in value_dict.items() ) else: return "(None)" @@ -275,9 +274,7 @@ def _terminate_one_execution(client, urn, cause, shouldPrint=True): client.terminate_execution(_identifier.WorkflowExecutionIdentifier.from_python_std(urn), cause) -def _update_one_launch_plan(urn, host, insecure, state): - client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure) - +def _update_one_launch_plan(client: _friendly_client.SynchronousFlyteClient, urn, state): if state == "active": state = _launch_plan.LaunchPlanState.ACTIVE else: @@ -306,6 +303,7 @@ def _render_schedule_expr(lp): _CONFIG_FLAGS = ["-c", "--config"] _PRINCIPAL_FLAGS = ["-r", "--principal"] _INSECURE_FLAGS = ["-i", "--insecure"] +_CERT_FLAGS = ["--cacert"] _project_option = _click.option(*_PROJECT_FLAGS, required=True, help="The project namespace to query.") _optional_project_option = _click.option( @@ -505,7 +503,7 @@ def make_context(self, cmd_name, args, parent=None): and param.name in parent.params and parent.params[param.name] is not None ): - prefix_args.extend([type(self)._PASSABLE_ARGS[param.name], _six.text_type(parent.params[param.name])]) + prefix_args.extend([type(self)._PASSABLE_ARGS[param.name], str(parent.params[param.name])]) # For flags, we don't append the value of the flag, otherwise click will fail to parse if param.name in type(self)._PASSABLE_FLAGS and param.name in parent.params and parent.params[param.name]: @@ -516,6 +514,8 @@ def make_context(self, cmd_name, args, parent=None): # flyte-cli setup-config -h localhost:30081 -i if cmd_name == "setup-config": ctx = super(_FlyteSubCommand, self).make_context(cmd_name, prefix_args + args, parent=parent) + ctx.obj = ctx.obj or {} + ctx.obj["cacert"] = parent.params["cacert"] or None return ctx config = parent.params["config"] @@ -562,8 +562,10 @@ def make_context(self, cmd_name, args, parent=None): # Use host url in config file if users don't specify the host url if _HOST_FLAGS[0] not in prefix_args: - prefix_args.extend([_HOST_FLAGS[0], _six.text_type(_HOST_URL)]) + prefix_args.extend([_HOST_FLAGS[0], str(_HOST_URL)]) ctx = super(_FlyteSubCommand, self).make_context(cmd_name, prefix_args + args, parent=parent) + ctx.obj = ctx.obj or {} + ctx.obj["cacert"] = parent.params["cacert"] or None return ctx @@ -607,10 +609,17 @@ def make_context(self, cmd_name, args, parent=None): help="[Optional] The name to pass to the sub-command (if applicable) If set again in the sub-command, " "the sub-command's parameter takes precedence.", ) +@_click.option( + *_CERT_FLAGS, + required=False, + type=str, + default=None, + help="[Optional] Path to certificate file to be used to do establish SSL connection with Admin", +) @_insecure_option @_click.group("flyte-cli", deprecated=True) @_click.pass_context -def _flyte_cli(ctx, host, config, project, domain, name, insecure): +def _flyte_cli(ctx, host, config, project, domain, name, cacert, insecure): """ Command line tool for interacting with all entities on the Flyte Platform. """ @@ -663,7 +672,8 @@ def list_task_names(project, domain, host, insecure, token, limit, show_all, sor a specific project and domain. """ _welcome_message() - client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure) + parent_ctx = _click.get_current_context(silent=True) + client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) _click.echo("Task Names Found in {}:{}\n".format(_tt(project), _tt(domain))) while True: @@ -705,7 +715,8 @@ def list_task_versions(project, domain, name, host, insecure, token, limit, show versions of that particular task (identifiable by {Project, Domain, Name}). """ _welcome_message() - client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure) + parent_ctx = _click.get_current_context(silent=True) + client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) _click.echo("Task Versions Found for {}:{}:{}\n".format(_tt(project), _tt(domain), _tt(name or "*"))) _click.echo("{:50} {:40}".format("Version", "Urn")) @@ -745,7 +756,8 @@ def get_task(urn, host, insecure): The URN of the versioned task is in the form of ``tsk::::``. """ _welcome_message() - client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure) + parent_ctx = _click.get_current_context(silent=True) + client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) t = client.get_task(_identifier.Identifier.from_python_std(urn)) _click.echo(_tt(t)) _click.echo("") @@ -784,7 +796,7 @@ def launch_task(project, domain, name, assumable_iam_role, kubernetes_service_ac text_args = _parse_args_into_dict(task_args) inputs = {} - for var_name, variable in _six.iteritems(task.interface.inputs): + for var_name, variable in task.interface.inputs.items(): sdk_type = _type_helpers.get_sdk_type_from_literal_type(variable.type) if var_name in text_args and text_args[var_name] is not None: inputs[var_name] = sdk_type.from_string(text_args[var_name]).to_python_std() @@ -818,7 +830,8 @@ def list_workflow_names(project, domain, host, insecure, token, limit, show_all, List the names of the workflows under a scope specified by ``{project, domain}``. """ _welcome_message() - client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure) + parent_ctx = _click.get_current_context(silent=True) + client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) _click.echo("Workflow Names Found in {}:{}\n".format(_tt(project), _tt(domain))) while True: @@ -860,7 +873,8 @@ def list_workflow_versions(project, domain, name, host, insecure, token, limit, versions of that particular workflow (identifiable by ``{project, domain, name}``). """ _welcome_message() - client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure) + parent_ctx = _click.get_current_context(silent=True) + client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) _click.echo("Workflow Versions Found for {}:{}:{}\n".format(_tt(project), _tt(domain), _tt(name or "*"))) _click.echo("{:50} {:40}".format("Version", "Urn")) @@ -900,7 +914,8 @@ def get_workflow(urn, host, insecure): ``wf::::`` """ _welcome_message() - client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure) + parent_ctx = _click.get_current_context(silent=True) + client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) _click.echo(client.get_workflow(_identifier.Identifier.from_python_std(urn))) # TODO: Print workflow pretty _click.echo("") @@ -927,7 +942,8 @@ def list_launch_plan_names(project, domain, host, insecure, token, limit, show_a List the names of the launch plans under the scope specified by {project, domain}. """ _welcome_message() - client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure) + parent_ctx = _click.get_current_context(silent=True) + client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) _click.echo("Launch Plan Names Found in {}:{}\n".format(_tt(project), _tt(domain))) while True: @@ -971,7 +987,8 @@ def list_active_launch_plans(project, domain, host, insecure, token, limit, show _click.echo("Active Launch Plan Found in {}:{}\n".format(_tt(project), _tt(domain))) _click.echo("{:30} {:50} {:80}".format("Schedule", "Version", "Urn")) - client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure) + parent_ctx = _click.get_current_context(silent=True) + client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) while True: active_lps, next_token = client.list_active_launch_plans_paginated( @@ -1040,7 +1057,8 @@ def list_launch_plan_versions( _click.echo("Launch Plan Versions Found for {}:{}:{}\n".format(_tt(project), _tt(domain), _tt(name))) _click.echo("{:50} {:80} {:30} {:15}".format("Version", "Urn", "Schedule", "Schedule State")) - client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure) + parent_ctx = _click.get_current_context(silent=True) + client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) while True: lp_list, next_token = client.list_launch_plans_paginated( @@ -1093,7 +1111,8 @@ def get_launch_plan(urn, host, insecure): The URN of a launch plan is in the form of ``lp::::`` """ _welcome_message() - client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure) + parent_ctx = _click.get_current_context(silent=True) + client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) _click.echo(_tt(client.get_launch_plan(_identifier.Identifier.from_python_std(urn)))) # TODO: Print launch plan pretty _click.echo("") @@ -1110,7 +1129,8 @@ def get_active_launch_plan(project, domain, name, host, insecure): List the versions of all the launch plans under the scope specified by {project, domain}. """ _welcome_message() - client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure) + parent_ctx = _click.get_current_context(silent=True) + client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) lp = client.get_active_launch_plan(_common_models.NamedEntityIdentifier(project, domain, name)) _click.echo("Active Launch Plan for {}:{}:{}\n".format(_tt(project), _tt(domain), _tt(name))) @@ -1125,13 +1145,15 @@ def get_active_launch_plan(project, domain, name, host, insecure): @_optional_urn_option def update_launch_plan(state, host, insecure, urn=None): _welcome_message() + parent_ctx = _click.get_current_context(silent=True) + client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) if urn is None: try: # Examine whether the input is from the named pipe if _stat.S_ISFIFO(_os.fstat(0).st_mode): for line in _sys.stdin.readlines(): - _update_one_launch_plan(urn=line.rstrip(), host=host, insecure=insecure, state=state) + _update_one_launch_plan(client, urn=line.rstrip(), state=state) else: # If the commandline parameter urn is not supplied, and neither # the input comes from a pipe, it means the user is not using @@ -1140,7 +1162,7 @@ def update_launch_plan(state, host, insecure, urn=None): except KeyboardInterrupt: _sys.stdout.flush() else: - _update_one_launch_plan(urn=urn, host=host, insecure=insecure, state=state) + _update_one_launch_plan(client, urn=urn, state=state) @_flyte_cli.command("execute-launch-plan", cls=_FlyteSubCommand) @@ -1206,8 +1228,8 @@ def watch_execution(host, insecure, urn): $ flyte-cli -h localhost:30081 watch-execution -u ex:flyteexamples:development:abc123 """ _welcome_message() - - client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure) + parent_ctx = _click.get_current_context(silent=True) + client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) ex_id = _identifier.WorkflowExecutionIdentifier.from_python_std(urn) execution = _workflow_execution_common.SdkWorkflowExecution.promote_from_model(client.get_execution(ex_id)) @@ -1252,7 +1274,8 @@ def relaunch_execution(project, domain, name, host, insecure, urn, principal, ve Users should use the get-execution and get-launch-plan commands to ascertain the names of inputs to use. """ _welcome_message() - client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure) + parent_ctx = _click.get_current_context(silent=True) + client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) _click.echo("Relaunching execution {}\n".format(_tt(urn))) existing_workflow_execution_identifier = _identifier.WorkflowExecutionIdentifier.from_python_std(urn) @@ -1270,7 +1293,7 @@ def relaunch_execution(project, domain, name, host, insecure, urn, principal, ve # Parse text inputs using the LP closure's parameter map to determine types. However, since all inputs are now # optional (because we can default to the original execution's), we reduce first to bare Variables. - variable_map = {k: v.var for k, v in _six.iteritems(expected_inputs.parameters)} + variable_map = {k: v.var for k, v in expected_inputs.parameters.items()} parsed_text_args = _parse_args_into_dict(lp_args) new_inputs = _construct_literal_map_from_variable_map(variable_map, parsed_text_args) if len(new_inputs.literals) > 0: @@ -1326,7 +1349,8 @@ def recover_execution(urn, name, host, insecure): Users should use the get-execution and get-launch-plan commands to ascertain the names of inputs to use. """ _welcome_message() - client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure) + parent_ctx = _click.get_current_context(silent=True) + client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) _click.echo("Recovering execution {}\n".format(_tt(urn))) @@ -1363,7 +1387,8 @@ def terminate_execution(host, insecure, cause, urn=None): -u lp:flyteexamples:development:some-execution:abc123 """ _welcome_message() - client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure) + parent_ctx = _click.get_current_context(silent=True) + client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) _click.echo("Killing the following executions:\n") _click.echo("{:100} {:40}".format("Urn", "Cause")) @@ -1415,7 +1440,8 @@ def list_executions(project, domain, host, insecure, token, limit, show_all, fil _click.echo("Executions Found in {}:{}\n".format(_tt(project), _tt(domain))) _click.echo("{:100} {:40} {:10}".format("Urn", "Name", "Status")) - client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure) + parent_ctx = _click.get_current_context(silent=True) + client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) while True: exec_ids, next_token = client.list_executions_paginated( @@ -1669,7 +1695,8 @@ def get_execution(urn, host, insecure, show_io, verbose): The URN of an execution is in the form of ``ex:::`` """ _welcome_message() - client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure) + parent_ctx = _click.get_current_context(silent=True) + client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) e = client.get_execution(_identifier.WorkflowExecutionIdentifier.from_python_std(urn)) node_execs = _get_all_node_executions(client, workflow_execution_identifier=e.id) _render_node_executions(client, node_execs, show_io, verbose, host, insecure, wf_execution=e) @@ -1683,7 +1710,8 @@ def get_execution(urn, host, insecure, show_io, verbose): @_verbose_option def get_child_executions(urn, host, insecure, show_io, verbose): _welcome_message() - client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure) + parent_ctx = _click.get_current_context(silent=True) + client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) node_execs = _get_all_node_executions( client, task_execution_identifier=_identifier.TaskExecutionIdentifier.from_python_std(urn), @@ -1703,7 +1731,8 @@ def register_project(identifier, name, description, host, insecure): """ _welcome_message() - client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure) + parent_ctx = _click.get_current_context(silent=True) + client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) client.register_project(_Project(identifier, name, description)) _click.echo("Registered project [id: {}, name: {}, description: {}]".format(identifier, name, description)) @@ -1722,7 +1751,8 @@ def list_projects(host, insecure, token, limit, show_all, filter, sort_by): """ _welcome_message() - client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure) + parent_ctx = _click.get_current_context(silent=True) + client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) _click.echo("Projects Found\n") while True: @@ -1755,7 +1785,9 @@ def archive_project(identifier, host, insecure): """ _welcome_message() - client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure) + parent_ctx = _click.get_current_context(silent=True) + client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) + client.update_project(_Project.archived_project(identifier)) _click.echo("Archived project [id: {}]".format(identifier)) @@ -1770,7 +1802,8 @@ def activate_project(identifier, host, insecure): """ _welcome_message() - client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure) + parent_ctx = _click.get_current_context(silent=True) + client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) client.update_project(_Project.active_project(identifier)) _click.echo("Activated project [id: {}]".format(identifier)) @@ -1877,16 +1910,13 @@ def patch_launch_plan(entity: _GeneratedProtocolMessageType) -> _GeneratedProtoc def _extract_and_register( - host: str, - insecure: bool, + client: _friendly_client.SynchronousFlyteClient, project: str, domain: str, version: str, file_paths: List[str], patches: Dict[int, Callable[[_GeneratedProtocolMessageType], _GeneratedProtocolMessageType]] = None, ): - client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure) - flyte_entities_list = _extract_files(project, domain, version, file_paths, patches) for id, flyte_entity in flyte_entities_list: _click.secho(f"Registering {id}", fg="yellow") @@ -1958,8 +1988,9 @@ def register_files( assumable_iam_role, kubernetes_service_account, output_location_prefix ) } - - _extract_and_register(host, insecure, project, domain, version, files, patches) + parent_ctx = _click.get_current_context(silent=True) + client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) + _extract_and_register(client, project, domain, version, files, patches) def _substitute_fast_register_task_args(args: List[str], full_remote_path: str, dest_dir: str) -> List[str]: @@ -2082,8 +2113,10 @@ def fast_register_task(entity: _GeneratedProtocolMessageType) -> _GeneratedProto assumable_iam_role, kubernetes_service_account, output_location_prefix ), } + parent_ctx = _click.get_current_context(silent=True) + client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) - _extract_and_register(host, insecure, project, domain, version, pb_files, patches) + _extract_and_register(client, project, domain, version, pb_files, patches) @_flyte_cli.command("update-workflow-meta", cls=_FlyteSubCommand) @@ -2099,7 +2132,8 @@ def update_workflow_meta(description, state, host, insecure, project, domain, na Updates a workflow entity under the scope specified by {project, domain, name} across versions. """ _welcome_message() - client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure) + parent_ctx = _click.get_current_context(silent=True) + client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) if state == "active": state = _named_entity.NamedEntityState.ACTIVE elif state == "archived": @@ -2124,7 +2158,8 @@ def update_task_meta(description, host, insecure, project, domain, name): Updates a task entity under the scope specified by {project, domain, name} across versions. """ _welcome_message() - client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure) + parent_ctx = _click.get_current_context(silent=True) + client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) client.update_named_entity( _core_identifier.ResourceType.TASK, _named_entity.NamedEntityIdentifier(project, domain, name), @@ -2145,7 +2180,8 @@ def update_launch_plan_meta(description, host, insecure, project, domain, name): Updates a launch plan entity under the scope specified by {project, domain, name} across versions. """ _welcome_message() - client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure) + parent_ctx = _click.get_current_context(silent=True) + client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) client.update_named_entity( _core_identifier.ResourceType.LAUNCH_PLAN, _named_entity.NamedEntityIdentifier(project, domain, name), @@ -2174,7 +2210,8 @@ def update_cluster_resource_attributes(host, insecure, project, domain, name, at --attributes projectQuotaCpu 1 --attributes projectQuotaMemory 500M """ _welcome_message() - client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure) + parent_ctx = _click.get_current_context(silent=True) + client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) cluster_resource_attributes = _ClusterResourceAttributes({attribute[0]: attribute[1] for attribute in attributes}) matching_attributes = _MatchingAttributes(cluster_resource_attributes=cluster_resource_attributes) @@ -2208,7 +2245,8 @@ def update_execution_queue_attributes(host, insecure, project, domain, name, tag --tags critical --tags gpu_intensive """ _welcome_message() - client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure) + parent_ctx = _click.get_current_context(silent=True) + client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) execution_queue_attributes = _ExecutionQueueAttributes(list(tags)) matching_attributes = _MatchingAttributes(execution_queue_attributes=execution_queue_attributes) @@ -2242,7 +2280,8 @@ def update_execution_cluster_label(host, insecure, project, domain, name, value) $ flyte-cli -h localhost:30081 -p flyteexamples -d development update-execution-cluster-label --value foo """ _welcome_message() - client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure) + parent_ctx = _click.get_current_context(silent=True) + client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) execution_cluster_label = _ExecutionClusterLabel(value) matching_attributes = _MatchingAttributes(execution_cluster_label=execution_cluster_label) @@ -2280,7 +2319,8 @@ def update_plugin_override(host, insecure, project, domain, name, task_type, plu --plugin-id my_cool_plugin --plugin-id my_fallback_plugin --missing-plugin-behavior FAIL """ _welcome_message() - client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure) + parent_ctx = _click.get_current_context(silent=True) + client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) plugin_override = _PluginOverride( task_type, list(plugin_id), _PluginOverride.string_to_enum(missing_plugin_behavior.upper()) ) @@ -2324,7 +2364,8 @@ def get_matching_attributes(host, insecure, project, domain, name, resource_type combination. """ _welcome_message() - client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure) + parent_ctx = _click.get_current_context(silent=True) + client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) if name is not None: attributes = client.get_workflow_attributes( @@ -2360,7 +2401,8 @@ def list_matching_attributes(host, insecure, resource_type): Fetches all matchable resources of the given resource type. """ _welcome_message() - client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure) + parent_ctx = _click.get_current_context(silent=True) + client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) attributes = client.list_matchable_attributes(_MatchableResource.string_to_enum(resource_type.upper())) for configuration in attributes.configurations: diff --git a/tests/flytekit/unit/remote/test_remote.py b/tests/flytekit/unit/remote/test_remote.py index 2bb8174069..5904d2efdf 100644 --- a/tests/flytekit/unit/remote/test_remote.py +++ b/tests/flytekit/unit/remote/test_remote.py @@ -179,4 +179,4 @@ def test_explicit_grpc_channel_credentials(mock_insecure, mock_url, mock_secure_ _ = FlyteRemote.from_config("project", "domain", grpc_credentials=credentials) assert mock_secure_channel.called assert mock_secure_channel.call_args[0][1] == credentials - assert not mock_ssl_channel_credentials.called + assert mock_ssl_channel_credentials.call_count == 1 From 7db24b83d8466810c45b9bc4340ef7216c9a3084 Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Mon, 20 Dec 2021 12:34:06 -0800 Subject: [PATCH 042/128] Add validation check to cacert switch (#787) Signed-off-by: Yee Hing Tong Signed-off-by: maximsmol --- flytekit/clis/flyte_cli/main.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/flytekit/clis/flyte_cli/main.py b/flytekit/clis/flyte_cli/main.py index 4dab451827..5686854fad 100644 --- a/flytekit/clis/flyte_cli/main.py +++ b/flytekit/clis/flyte_cli/main.py @@ -623,6 +623,8 @@ def _flyte_cli(ctx, host, config, project, domain, name, cacert, insecure): """ Command line tool for interacting with all entities on the Flyte Platform. """ + if cacert and insecure: + raise _user_exceptions.FlyteValidationException(f"Should not pass both certificate and insecure options!") ######################################################################################################################## From 718007538bad6d0441897a302260a7c9e193772c Mon Sep 17 00:00:00 2001 From: Eduardo Apolinario <653394+eapolinario@users.noreply.github.com> Date: Mon, 20 Dec 2021 12:41:55 -0800 Subject: [PATCH 043/128] Remove pytz constraint (#786) * Remove pytz constraint from setup.py Signed-off-by: Eduardo Apolinario * Regenerate requirements files Signed-off-by: Eduardo Apolinario * Put pytz back Signed-off-by: Eduardo Apolinario * make requirements.txt Signed-off-by: Eduardo Apolinario Co-authored-by: Eduardo Apolinario Signed-off-by: maximsmol --- dev-requirements.txt | 29 ++++++----- doc-requirements.txt | 49 +++++++++---------- requirements-spark2.txt | 44 ++++++++--------- requirements.txt | 40 +++++++-------- setup.py | 2 +- .../workflows/requirements.txt | 35 +++++-------- 6 files changed, 90 insertions(+), 109 deletions(-) diff --git a/dev-requirements.txt b/dev-requirements.txt index 4ba14b6a8c..cfea1f15a8 100644 --- a/dev-requirements.txt +++ b/dev-requirements.txt @@ -44,7 +44,7 @@ chardet==4.0.0 # via # -c requirements.txt # binaryornot -charset-normalizer==2.0.8 +charset-normalizer==2.0.9 # via # -c requirements.txt # requests @@ -69,11 +69,11 @@ cookiecutter==1.7.3 # flytekit coverage[toml]==6.2 # via -r dev-requirements.in -croniter==1.0.15 +croniter==1.1.0 # via # -c requirements.txt # flytekit -cryptography==36.0.0 +cryptography==36.0.1 # via # -c requirements.txt # paramiko @@ -94,7 +94,7 @@ diskcache==5.3.0 # via # -c requirements.txt # flytekit -distlib==0.3.3 +distlib==0.3.4 # via virtualenv distro==1.6.0 # via docker-compose @@ -118,11 +118,11 @@ docstring-parser==0.13 # flytekit filelock==3.4.0 # via virtualenv -flyteidl==0.21.11 +flyteidl==0.21.13 # via # -c requirements.txt # flytekit -grpcio==1.42.0 +grpcio==1.43.0 # via # -c requirements.txt # flytekit @@ -132,7 +132,7 @@ idna==3.3 # via # -c requirements.txt # requests -importlib-metadata==4.8.2 +importlib-metadata==4.10.0 # via # -c requirements.txt # keyring @@ -183,20 +183,20 @@ marshmallow-jsonschema==0.13.0 # flytekit mock==4.0.3 # via -r dev-requirements.in -mypy==0.910 +mypy==0.920 # via -r dev-requirements.in mypy-extensions==0.4.3 # via # -c requirements.txt # mypy # typing-inspect -natsort==8.0.0 +natsort==8.0.2 # via # -c requirements.txt # flytekit nodeenv==1.6.0 # via pre-commit -numpy==1.21.4 +numpy==1.21.5 # via # -c requirements.txt # pandas @@ -205,7 +205,7 @@ packaging==21.3 # via # -c requirements.txt # pytest -pandas==1.3.4 +pandas==1.3.5 # via # -c requirements.txt # flytekit @@ -285,10 +285,9 @@ pytimeparse==1.1.8 # via # -c requirements.txt # flytekit -pytz==2018.4 +pytz==2021.3 # via # -c requirements.txt - # flytekit # pandas pyyaml==5.4.1 # via @@ -349,13 +348,13 @@ texttable==1.6.4 # via docker-compose toml==0.10.2 # via - # mypy # pre-commit # pytest -tomli==1.2.2 +tomli==1.2.3 # via # -c requirements.txt # coverage + # mypy typing-extensions==4.0.1 # via # -c requirements.txt diff --git a/doc-requirements.txt b/doc-requirements.txt index 1913e64cf2..df16bb0ff4 100644 --- a/doc-requirements.txt +++ b/doc-requirements.txt @@ -29,13 +29,13 @@ beautifulsoup4==4.10.0 # sphinx-material binaryornot==0.4.4 # via cookiecutter -black==21.11b1 +black==21.12b0 # via papermill bleach==4.1.0 # via nbconvert -boto3==1.20.17 +boto3==1.20.24 # via sagemaker-training -botocore==1.23.17 +botocore==1.23.24 # via # boto3 # s3transfer @@ -48,7 +48,7 @@ cffi==1.15.0 # pynacl chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.8 +charset-normalizer==2.0.9 # via requests checksumdir==1.2.0 # via flytekit @@ -63,9 +63,9 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.0.15 +croniter==1.1.0 # via flytekit -cryptography==36.0.0 +cryptography==36.0.1 # via # -r doc-requirements.in # paramiko @@ -95,15 +95,15 @@ entrypoints==0.3 # jupyter-client # nbconvert # papermill -flyteidl==0.21.11 +flyteidl==0.21.13 # via flytekit furo @ git+git://github.com/flyteorg/furo@main # via -r doc-requirements.in -gevent==21.8.0 +gevent==21.12.0 # via sagemaker-training greenlet==1.1.2 # via gevent -grpcio==1.42.0 +grpcio==1.43.0 # via # -r doc-requirements.in # flytekit @@ -113,13 +113,13 @@ idna==3.3 # via requests imagesize==1.3.0 # via sphinx -importlib-metadata==4.8.2 +importlib-metadata==4.10.0 # via keyring inotify_simple==1.2.1 # via sagemaker-training ipykernel==5.5.6 # via flytekit -ipython==7.30.0 +ipython==7.30.1 # via ipykernel ipython-genutils==0.2.0 # via @@ -144,7 +144,7 @@ jmespath==0.10.0 # via # boto3 # botocore -jsonschema==4.2.1 +jsonschema==4.3.2 # via nbformat jupyter-client==7.1.0 # via @@ -161,9 +161,9 @@ k8s-proto==0.0.3 # via flytekit keyring==23.4.0 # via flytekit -lazy-object-proxy==1.6.0 +lazy-object-proxy==1.7.1 # via astroid -lxml==4.6.5 +lxml==4.7.1 # via sphinx-material markupsafe==2.0.1 # via jinja2 @@ -184,7 +184,7 @@ mypy-extensions==0.4.3 # via # black # typing-inspect -natsort==8.0.0 +natsort==8.0.2 # via flytekit nbclient==0.5.9 # via @@ -197,11 +197,11 @@ nbformat==5.1.3 # nbclient # nbconvert # papermill -nest-asyncio==1.5.2 +nest-asyncio==1.5.4 # via # jupyter-client # nbclient -numpy==1.21.4 +numpy==1.21.5 # via # flytekit # pandas @@ -212,7 +212,7 @@ packaging==21.3 # via # bleach # sphinx -pandas==1.3.4 +pandas==1.3.5 # via flytekit pandocfilters==1.5.0 # via nbconvert @@ -232,7 +232,7 @@ platformdirs==2.4.0 # via black poyo==0.5.0 # via cookiecutter -prompt-toolkit==3.0.23 +prompt-toolkit==3.0.24 # via ipython protobuf==3.19.1 # via @@ -283,10 +283,9 @@ python-slugify[unidecode]==5.0.2 # sphinx-material pytimeparse==1.1.8 # via flytekit -pytz==2018.4 +pytz==2021.3 # via # babel - # flytekit # pandas pyyaml==6.0 # via @@ -295,9 +294,7 @@ pyyaml==6.0 pyzmq==22.3.0 # via jupyter-client regex==2021.11.10 - # via - # black - # docker-image-py + # via docker-image-py requests==2.26.0 # via # cookiecutter @@ -339,7 +336,7 @@ sortedcontainers==2.4.0 # via flytekit soupsieve==2.3.1 # via beautifulsoup4 -sphinx==4.3.1 +sphinx==4.3.2 # via # -r doc-requirements.in # furo @@ -391,7 +388,7 @@ textwrap3==0.9.2 # via ansiwrap thrift==0.15.0 # via hmsclient -tomli==1.2.2 +tomli==1.2.3 # via black tornado==6.1 # via diff --git a/requirements-spark2.txt b/requirements-spark2.txt index a0457279bc..0ba0d63ded 100644 --- a/requirements-spark2.txt +++ b/requirements-spark2.txt @@ -22,13 +22,13 @@ bcrypt==3.2.0 # via paramiko binaryornot==0.4.4 # via cookiecutter -black==21.11b1 +black==21.12b0 # via papermill bleach==4.1.0 # via nbconvert -boto3==1.20.17 +boto3==1.20.24 # via sagemaker-training -botocore==1.23.17 +botocore==1.23.24 # via # boto3 # s3transfer @@ -41,7 +41,7 @@ cffi==1.15.0 # pynacl chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.8 +charset-normalizer==2.0.9 # via requests checksumdir==1.2.0 # via flytekit @@ -56,9 +56,9 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.0.15 +croniter==1.1.0 # via flytekit -cryptography==36.0.0 +cryptography==36.0.1 # via # paramiko # secretstorage @@ -83,25 +83,25 @@ entrypoints==0.3 # jupyter-client # nbconvert # papermill -flyteidl==0.21.11 +flyteidl==0.21.13 # via flytekit -gevent==21.8.0 +gevent==21.12.0 # via sagemaker-training greenlet==1.1.2 # via gevent -grpcio==1.42.0 +grpcio==1.43.0 # via flytekit hmsclient==0.1.1 # via flytekit idna==3.3 # via requests -importlib-metadata==4.8.2 +importlib-metadata==4.10.0 # via keyring inotify_simple==1.2.1 # via sagemaker-training ipykernel==5.5.6 # via flytekit -ipython==7.30.0 +ipython==7.30.1 # via ipykernel ipython-genutils==0.2.0 # via @@ -162,7 +162,7 @@ mypy-extensions==0.4.3 # via # black # typing-inspect -natsort==8.0.0 +natsort==8.0.2 # via flytekit nbclient==0.5.9 # via @@ -175,11 +175,11 @@ nbformat==5.1.3 # nbclient # nbconvert # papermill -nest-asyncio==1.5.2 +nest-asyncio==1.5.4 # via # jupyter-client # nbclient -numpy==1.21.4 +numpy==1.21.5 # via # flytekit # pandas @@ -188,7 +188,7 @@ numpy==1.21.4 # scipy packaging==21.3 # via bleach -pandas==1.3.4 +pandas==1.3.5 # via flytekit pandocfilters==1.5.0 # via nbconvert @@ -208,7 +208,7 @@ platformdirs==2.4.0 # via black poyo==0.5.0 # via cookiecutter -prompt-toolkit==3.0.23 +prompt-toolkit==3.0.24 # via ipython protobuf==3.19.1 # via @@ -255,10 +255,8 @@ python-slugify==5.0.2 # via cookiecutter pytimeparse==1.1.8 # via flytekit -pytz==2018.4 - # via - # flytekit - # pandas +pytz==2021.3 + # via pandas pyyaml==5.4.1 # via # -r requirements.in @@ -266,9 +264,7 @@ pyyaml==5.4.1 pyzmq==22.3.0 # via jupyter-client regex==2021.11.10 - # via - # black - # docker-image-py + # via docker-image-py requests==2.26.0 # via # cookiecutter @@ -317,7 +313,7 @@ textwrap3==0.9.2 # via ansiwrap thrift==0.15.0 # via hmsclient -tomli==1.2.2 +tomli==1.2.3 # via black tornado==6.1 # via diff --git a/requirements.txt b/requirements.txt index 344af69890..9aacef2d24 100644 --- a/requirements.txt +++ b/requirements.txt @@ -20,13 +20,13 @@ bcrypt==3.2.0 # via paramiko binaryornot==0.4.4 # via cookiecutter -black==21.11b1 +black==21.12b0 # via papermill bleach==4.1.0 # via nbconvert -boto3==1.20.17 +boto3==1.20.25 # via sagemaker-training -botocore==1.23.17 +botocore==1.23.25 # via # boto3 # s3transfer @@ -39,7 +39,7 @@ cffi==1.15.0 # pynacl chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.8 +charset-normalizer==2.0.9 # via requests checksumdir==1.2.0 # via flytekit @@ -54,9 +54,9 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.0.15 +croniter==1.1.0 # via flytekit -cryptography==36.0.0 +cryptography==36.0.1 # via # paramiko # secretstorage @@ -81,25 +81,25 @@ entrypoints==0.3 # jupyter-client # nbconvert # papermill -flyteidl==0.21.11 +flyteidl==0.21.13 # via flytekit -gevent==21.8.0 +gevent==21.12.0 # via sagemaker-training greenlet==1.1.2 # via gevent -grpcio==1.42.0 +grpcio==1.43.0 # via flytekit hmsclient==0.1.1 # via flytekit idna==3.3 # via requests -importlib-metadata==4.8.2 +importlib-metadata==4.10.0 # via keyring inotify_simple==1.2.1 # via sagemaker-training ipykernel==5.5.6 # via flytekit -ipython==7.30.0 +ipython==7.30.1 # via ipykernel ipython-genutils==0.2.0 # via @@ -160,7 +160,7 @@ mypy-extensions==0.4.3 # via # black # typing-inspect -natsort==8.0.0 +natsort==8.0.2 # via flytekit nbclient==0.5.9 # via @@ -173,11 +173,11 @@ nbformat==5.1.3 # nbclient # nbconvert # papermill -nest-asyncio==1.5.2 +nest-asyncio==1.5.4 # via # jupyter-client # nbclient -numpy==1.21.4 +numpy==1.21.5 # via # flytekit # pandas @@ -186,7 +186,7 @@ numpy==1.21.4 # scipy packaging==21.3 # via bleach -pandas==1.3.4 +pandas==1.3.5 # via flytekit pandocfilters==1.5.0 # via nbconvert @@ -206,7 +206,7 @@ platformdirs==2.4.0 # via black poyo==0.5.0 # via cookiecutter -prompt-toolkit==3.0.23 +prompt-toolkit==3.0.24 # via ipython protobuf==3.19.1 # via @@ -253,7 +253,7 @@ python-slugify==5.0.2 # via cookiecutter pytimeparse==1.1.8 # via flytekit -pytz==2018.4 +pytz==2021.3 # via # flytekit # pandas @@ -264,9 +264,7 @@ pyyaml==5.4.1 pyzmq==22.3.0 # via jupyter-client regex==2021.11.10 - # via - # black - # docker-image-py + # via docker-image-py requests==2.26.0 # via # cookiecutter @@ -315,7 +313,7 @@ textwrap3==0.9.2 # via ansiwrap thrift==0.15.0 # via hmsclient -tomli==1.2.2 +tomli==1.2.3 # via black tornado==6.1 # via diff --git a/setup.py b/setup.py index 305e3d3700..9d1a1acbc7 100644 --- a/setup.py +++ b/setup.py @@ -76,7 +76,7 @@ "protobuf>=3.6.1,<4", "python-json-logger>=2.0.0", "pytimeparse>=1.1.8,<2.0.0", - "pytz>=2017.2,<2018.5", + "pytz", "keyring>=18.0.1", "requests>=2.18.4,<3.0.0", "responses>=0.10.7", diff --git a/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt b/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt index f9ecc1e38d..cec1da7299 100644 --- a/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt +++ b/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt @@ -14,7 +14,7 @@ cffi==1.15.0 # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.8 +charset-normalizer==2.0.9 # via requests checksumdir==1.2.0 # via flytekit @@ -26,9 +26,9 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.0.15 +croniter==1.1.0 # via flytekit -cryptography==36.0.0 +cryptography==36.0.1 # via secretstorage cycler==0.11.0 # via matplotlib @@ -44,17 +44,17 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.21.11 +flyteidl==0.21.13 # via flytekit -flytekit==0.24.0 +flytekit==0.25.0 # via -r tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.in -fonttools==4.28.2 +fonttools==4.28.5 # via matplotlib -grpcio==1.42.0 +grpcio==1.43.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.8.2 +importlib-metadata==4.10.0 # via keyring jeepney==0.7.1 # via @@ -83,13 +83,13 @@ marshmallow-enum==1.5.1 # via dataclasses-json marshmallow-jsonschema==0.13.0 # via flytekit -matplotlib==3.5.0 +matplotlib==3.5.1 # via -r tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.in mypy-extensions==0.4.3 # via typing-inspect -natsort==8.0.0 +natsort==8.0.2 # via flytekit -numpy==1.21.4 +numpy==1.21.5 # via # matplotlib # opencv-python @@ -98,10 +98,8 @@ numpy==1.21.4 opencv-python==4.5.4.60 # via -r tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.in packaging==21.3 - # via - # matplotlib - # setuptools-scm -pandas==1.3.4 + # via matplotlib +pandas==1.3.5 # via flytekit pillow==8.4.0 # via matplotlib @@ -151,8 +149,6 @@ retry==0.9.2 # via flytekit secretstorage==3.3.1 # via keyring -setuptools-scm==6.3.2 - # via matplotlib six==1.16.0 # via # cookiecutter @@ -166,8 +162,6 @@ statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -tomli==1.2.2 - # via setuptools-scm typing-extensions==4.0.1 # via typing-inspect typing-inspect==0.7.1 @@ -187,6 +181,3 @@ wrapt==1.13.3 # flytekit zipp==3.6.0 # via importlib-metadata - -# The following packages are considered to be unsafe in a requirements file: -# setuptools From 50da846f94c7285a3733381712e88684371f1976 Mon Sep 17 00:00:00 2001 From: Eduardo Apolinario <653394+eapolinario@users.noreply.github.com> Date: Mon, 20 Dec 2021 14:38:21 -0800 Subject: [PATCH 044/128] Lint: remove f-string misuse (#788) Signed-off-by: Eduardo Apolinario Signed-off-by: maximsmol --- flytekit/clis/flyte_cli/main.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flytekit/clis/flyte_cli/main.py b/flytekit/clis/flyte_cli/main.py index 5686854fad..438e91141f 100644 --- a/flytekit/clis/flyte_cli/main.py +++ b/flytekit/clis/flyte_cli/main.py @@ -624,7 +624,7 @@ def _flyte_cli(ctx, host, config, project, domain, name, cacert, insecure): Command line tool for interacting with all entities on the Flyte Platform. """ if cacert and insecure: - raise _user_exceptions.FlyteValidationException(f"Should not pass both certificate and insecure options!") + raise _user_exceptions.FlyteValidationException("Should not pass both certificate and insecure options!") ######################################################################################################################## From 46c692e77343879f6d0e8c7108774f005ce33ae2 Mon Sep 17 00:00:00 2001 From: Eduardo Apolinario <653394+eapolinario@users.noreply.github.com> Date: Thu, 23 Dec 2021 10:52:49 -0800 Subject: [PATCH 045/128] Pyarrow greater than 4.0.0 (#790) * pyarrow>=4.0.0 Signed-off-by: Eduardo Apolinario * Regenerate requirements. Signed-off-by: Eduardo Apolinario Co-authored-by: Eduardo Apolinario Signed-off-by: maximsmol --- dev-requirements.txt | 5 +++-- doc-requirements.txt | 7 ++++--- requirements-spark2.txt | 10 ++++++---- requirements.txt | 6 +++--- setup.py | 4 ++-- .../remote/mock_flyte_repo/workflows/requirements.txt | 2 +- 6 files changed, 19 insertions(+), 15 deletions(-) diff --git a/dev-requirements.txt b/dev-requirements.txt index cfea1f15a8..8cc4f7ba3f 100644 --- a/dev-requirements.txt +++ b/dev-requirements.txt @@ -183,7 +183,7 @@ marshmallow-jsonschema==0.13.0 # flytekit mock==4.0.3 # via -r dev-requirements.in -mypy==0.920 +mypy==0.930 # via -r dev-requirements.in mypy-extensions==0.4.3 # via @@ -288,6 +288,7 @@ pytimeparse==1.1.8 pytz==2021.3 # via # -c requirements.txt + # flytekit # pandas pyyaml==5.4.1 # via @@ -376,7 +377,7 @@ websocket-client==0.59.0 # via # docker # docker-compose -wheel==0.37.0 +wheel==0.37.1 # via # -c requirements.txt # flytekit diff --git a/doc-requirements.txt b/doc-requirements.txt index df16bb0ff4..cba57b13b7 100644 --- a/doc-requirements.txt +++ b/doc-requirements.txt @@ -33,9 +33,9 @@ black==21.12b0 # via papermill bleach==4.1.0 # via nbconvert -boto3==1.20.24 +boto3==1.20.26 # via sagemaker-training -botocore==1.23.24 +botocore==1.23.26 # via # boto3 # s3transfer @@ -286,6 +286,7 @@ pytimeparse==1.1.8 pytz==2021.3 # via # babel + # flytekit # pandas pyyaml==6.0 # via @@ -429,7 +430,7 @@ webencodings==0.5.1 # via bleach werkzeug==2.0.2 # via sagemaker-training -wheel==0.37.0 +wheel==0.37.1 # via flytekit wrapt==1.13.3 # via diff --git a/requirements-spark2.txt b/requirements-spark2.txt index 0ba0d63ded..a48ec4b5e2 100644 --- a/requirements-spark2.txt +++ b/requirements-spark2.txt @@ -26,9 +26,9 @@ black==21.12b0 # via papermill bleach==4.1.0 # via nbconvert -boto3==1.20.24 +boto3==1.20.26 # via sagemaker-training -botocore==1.23.24 +botocore==1.23.26 # via # boto3 # s3transfer @@ -256,7 +256,9 @@ python-slugify==5.0.2 pytimeparse==1.1.8 # via flytekit pytz==2021.3 - # via pandas + # via + # flytekit + # pandas pyyaml==5.4.1 # via # -r requirements.in @@ -349,7 +351,7 @@ webencodings==0.5.1 # via bleach werkzeug==2.0.2 # via sagemaker-training -wheel==0.37.0 +wheel==0.37.1 # via flytekit wrapt==1.13.3 # via diff --git a/requirements.txt b/requirements.txt index 9aacef2d24..04cdd38a16 100644 --- a/requirements.txt +++ b/requirements.txt @@ -24,9 +24,9 @@ black==21.12b0 # via papermill bleach==4.1.0 # via nbconvert -boto3==1.20.25 +boto3==1.20.26 # via sagemaker-training -botocore==1.23.25 +botocore==1.23.26 # via # boto3 # s3transfer @@ -349,7 +349,7 @@ webencodings==0.5.1 # via bleach werkzeug==2.0.2 # via sagemaker-training -wheel==0.37.0 +wheel==0.37.1 # via flytekit wrapt==1.13.3 # via diff --git a/setup.py b/setup.py index 9d1a1acbc7..f781690027 100644 --- a/setup.py +++ b/setup.py @@ -23,7 +23,7 @@ spark = ["pyspark>=2.4.0,<3.0.0"] spark3 = ["pyspark>=3.0.0"] sidecar = ["k8s-proto>=0.0.3,<1.0.0"] -schema = ["numpy>=1.14.0,<2.0.0", "pandas>=0.22.0,<2.0.0", "pyarrow>=6.0.0"] +schema = ["numpy>=1.14.0,<2.0.0", "pandas>=0.22.0,<2.0.0", "pyarrow>=4.0.0"] hive_sensor = ["hmsclient>=0.0.1,<1.0.0"] notebook = ["papermill>=1.2.0", "nbconvert>=6.0.7", "ipykernel>=5.0.0,<6.0.0"] sagemaker = ["sagemaker-training>=3.6.2,<4.0.0"] @@ -67,7 +67,7 @@ "flyteidl>=0.21.4", "wheel>=0.30.0,<1.0.0", "pandas>=1.0.0,<2.0.0", - "pyarrow>=6.0.0,<7.0.0", + "pyarrow>=4.0.0,<7.0.0", "click>=6.6,<8.0", "croniter>=0.3.20,<4.0.0", "deprecated>=1.0,<2.0", diff --git a/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt b/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt index cec1da7299..627e87980e 100644 --- a/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt +++ b/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt @@ -171,7 +171,7 @@ urllib3==1.26.7 # flytekit # requests # responses -wheel==0.37.0 +wheel==0.37.1 # via # -r tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.in # flytekit From 46fb085481b56bdccc442ad6ddaf29ee049cc961 Mon Sep 17 00:00:00 2001 From: Eduardo Apolinario <653394+eapolinario@users.noreply.github.com> Date: Tue, 28 Dec 2021 14:00:24 -0800 Subject: [PATCH 046/128] Support python 3.10 (#791) * [wip] Support python 3.10 Signed-off-by: Eduardo Apolinario * Add Dockerfile.py310 Signed-off-by: Eduardo Apolinario * Stringify python version Signed-off-by: Eduardo Apolinario * Skip flytekit-modin plugin tests on 3.10 Signed-off-by: Eduardo Apolinario * Add 3.9 and 3.10 to list of supported version in plugins Signed-off-by: Eduardo Apolinario * Comment why flytekit-modin is not running on 3.10 and disable fail-fast in plugin tests Signed-off-by: Eduardo Apolinario Co-authored-by: Eduardo Apolinario --- .github/workflows/pythonbuild.yml | 12 +++++- .github/workflows/pythonpublish.yml | 43 ++++++--------------- Dockerfile.py310 | 16 ++++++++ plugins/flytekit-aws-athena/setup.py | 2 + plugins/flytekit-aws-sagemaker/setup.py | 2 + plugins/flytekit-data-fsspec/setup.py | 2 + plugins/flytekit-dolt/setup.py | 2 + plugins/flytekit-greatexpectations/setup.py | 2 + plugins/flytekit-hive/setup.py | 2 + plugins/flytekit-k8s-pod/setup.py | 2 + plugins/flytekit-kf-mpi/setup.py | 2 + plugins/flytekit-kf-pytorch/setup.py | 2 + plugins/flytekit-kf-tensorflow/setup.py | 2 + plugins/flytekit-pandera/setup.py | 2 + plugins/flytekit-papermill/setup.py | 2 + plugins/flytekit-snowflake/setup.py | 2 + plugins/flytekit-spark/setup.py | 2 + plugins/flytekit-sqlalchemy/setup.py | 2 + setup.py | 1 + 19 files changed, 68 insertions(+), 34 deletions(-) create mode 100644 Dockerfile.py310 diff --git a/.github/workflows/pythonbuild.yml b/.github/workflows/pythonbuild.yml index ba633c0a19..59085ec468 100644 --- a/.github/workflows/pythonbuild.yml +++ b/.github/workflows/pythonbuild.yml @@ -12,13 +12,15 @@ jobs: strategy: fail-fast: false matrix: - python-version: [3.7, 3.8, 3.9] + python-version: ["3.7", "3.8", "3.9", "3.10"] spark-version-suffix: ["", "-spark2"] exclude: - python-version: 3.8 spark-version-suffix: "-spark2" - python-version: 3.9 spark-version-suffix: "-spark2" + - python-version: 3.10 + spark-version-suffix: "-spark2" steps: - uses: actions/checkout@v2 - name: Set up Python ${{ matrix.python-version }} @@ -55,8 +57,9 @@ jobs: build-plugins: runs-on: ubuntu-latest strategy: + fail-fast: false matrix: - python-version: [3.8, 3.9] + python-version: ["3.8", "3.9", "3.10"] plugin-names: - flytekit-aws-athena - flytekit-aws-sagemaker @@ -74,6 +77,11 @@ jobs: - flytekit-pandera - flytekit-snowflake - flytekit-modin + exclude: + # flytekit-modin depends on ray which does not have a 3.10 wheel yet. + # Issue tracked in https://github.com/ray-project/ray/issues/19116. + - python-version: 3.10 + plugin-names: "flytekit-modin" steps: - uses: actions/checkout@v2 - name: Set up Python ${{ matrix.python-version }} diff --git a/.github/workflows/pythonpublish.yml b/.github/workflows/pythonpublish.yml index d203a75ff1..0c64124273 100644 --- a/.github/workflows/pythonpublish.yml +++ b/.github/workflows/pythonpublish.yml @@ -52,45 +52,24 @@ jobs: uses: jakejarvis/wait-action@master with: time: '180s' - - name: Build & Push Flytekit Python3.7 Docker Image to Github Registry - uses: whoan/docker-build-with-cache-action@v5 - with: - # https://docs.github.com/en/packages/learn-github-packages/publishing-a-package - username: "${{ secrets.FLYTE_BOT_USERNAME }}" - password: "${{ secrets.FLYTE_BOT_PAT }}" - image_name: ${{ github.repository_owner }}/flytekit - image_tag: py37-latest,py37-${{ github.sha }},py37-${{ steps.bump.outputs.version }} - push_git_tag: true - push_image_and_stages: true - registry: ghcr.io - build_extra_args: "--compress=true --build-arg=VERSION=${{ steps.bump.outputs.version }} --build-arg=DOCKER_IMAGE=ghcr.io/flyteorg/flytekit:py37-${{ steps.bump.outputs.version }}" - context: . - dockerfile: Dockerfile.py37 - - name: Build & Push Flytekit Python3.8 Docker Image to Github Registry - uses: whoan/docker-build-with-cache-action@v5 - with: - # https://docs.github.com/en/packages/learn-github-packages/publishing-a-package - username: "${{ secrets.FLYTE_BOT_USERNAME }}" - password: "${{ secrets.FLYTE_BOT_PAT }}" - image_name: ${{ github.repository_owner }}/flytekit - image_tag: py38-latest,py38-${{ github.sha }},py38-${{ steps.bump.outputs.version }} - push_git_tag: true - push_image_and_stages: true - registry: ghcr.io - build_extra_args: "--compress=true --build-arg=VERSION=${{ steps.bump.outputs.version }} --build-arg=DOCKER_IMAGE=ghcr.io/flyteorg/flytekit:py38-${{ steps.bump.outputs.version }}" - context: . - dockerfile: Dockerfile.py38 - - name: Build & Push Flytekit Python3.9 Docker Image to Github Registry + + build-and-push-docker-images: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.7", "3.8", "3.9", "3.10"] + steps: + - name: Build & Push Flytekit Python${{ matrix.python-version }} Docker Image to Github Registry uses: whoan/docker-build-with-cache-action@v5 with: # https://docs.github.com/en/packages/learn-github-packages/publishing-a-package username: "${{ secrets.FLYTE_BOT_USERNAME }}" password: "${{ secrets.FLYTE_BOT_PAT }}" image_name: ${{ github.repository_owner }}/flytekit - image_tag: py39-latest,py39-${{ github.sha }},py39-${{ steps.bump.outputs.version }} + image_tag: py${{ matrix.python-version }}-latest,py${{ matrix.python-version }}-${{ github.sha }},py${{ matrix.python-version }}-${{ steps.bump.outputs.version }} push_git_tag: true push_image_and_stages: true registry: ghcr.io - build_extra_args: "--compress=true --build-arg=VERSION=${{ steps.bump.outputs.version }} --build-arg=DOCKER_IMAGE=ghcr.io/flyteorg/flytekit:py39-${{ steps.bump.outputs.version }}" + build_extra_args: "--compress=true --build-arg=VERSION=${{ steps.bump.outputs.version }} --build-arg=DOCKER_IMAGE=ghcr.io/flyteorg/flytekit:py${{ matrix.python-version }}-${{ steps.bump.outputs.version }}" context: . - dockerfile: Dockerfile.py39 + dockerfile: Dockerfile.py${{ matrix.python-version }} diff --git a/Dockerfile.py310 b/Dockerfile.py310 new file mode 100644 index 0000000000..1dc6b117b0 --- /dev/null +++ b/Dockerfile.py310 @@ -0,0 +1,16 @@ +FROM python:3.10-slim-buster + +MAINTAINER Flyte Team +LABEL org.opencontainers.image.source https://github.com/flyteorg/flytekit + +RUN pip install awscli +RUN pip install gsutil + +ARG VERSION +ARG DOCKER_IMAGE + +RUN pip install -U flytekit==$VERSION + +WORKDIR /app + +ENV FLYTE_INTERNAL_IMAGE "$DOCKER_IMAGE" diff --git a/plugins/flytekit-aws-athena/setup.py b/plugins/flytekit-aws-athena/setup.py index a179663fe8..298c3a5547 100644 --- a/plugins/flytekit-aws-athena/setup.py +++ b/plugins/flytekit-aws-athena/setup.py @@ -25,6 +25,8 @@ "License :: OSI Approved :: Apache Software License", "Programming Language :: Python :: 3.7", "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", "Topic :: Scientific/Engineering", "Topic :: Scientific/Engineering :: Artificial Intelligence", "Topic :: Software Development", diff --git a/plugins/flytekit-aws-sagemaker/setup.py b/plugins/flytekit-aws-sagemaker/setup.py index 7db7e813e4..1b171e631e 100644 --- a/plugins/flytekit-aws-sagemaker/setup.py +++ b/plugins/flytekit-aws-sagemaker/setup.py @@ -26,6 +26,8 @@ "License :: OSI Approved :: Apache Software License", "Programming Language :: Python :: 3.7", "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", "Topic :: Scientific/Engineering", "Topic :: Scientific/Engineering :: Artificial Intelligence", "Topic :: Software Development", diff --git a/plugins/flytekit-data-fsspec/setup.py b/plugins/flytekit-data-fsspec/setup.py index a0870aea76..3830d9d168 100644 --- a/plugins/flytekit-data-fsspec/setup.py +++ b/plugins/flytekit-data-fsspec/setup.py @@ -28,6 +28,8 @@ "License :: OSI Approved :: Apache Software License", "Programming Language :: Python :: 3.7", "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", "Topic :: Scientific/Engineering", "Topic :: Scientific/Engineering :: Artificial Intelligence", "Topic :: Software Development", diff --git a/plugins/flytekit-dolt/setup.py b/plugins/flytekit-dolt/setup.py index 5f808b480c..d950bdf592 100644 --- a/plugins/flytekit-dolt/setup.py +++ b/plugins/flytekit-dolt/setup.py @@ -41,6 +41,8 @@ def run(self): "License :: OSI Approved :: Apache Software License", "Programming Language :: Python :: 3.7", "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", "Topic :: Scientific/Engineering", "Topic :: Scientific/Engineering :: Artificial Intelligence", "Topic :: Software Development", diff --git a/plugins/flytekit-greatexpectations/setup.py b/plugins/flytekit-greatexpectations/setup.py index 2716d86b8a..3541415664 100644 --- a/plugins/flytekit-greatexpectations/setup.py +++ b/plugins/flytekit-greatexpectations/setup.py @@ -25,6 +25,8 @@ "License :: OSI Approved :: Apache Software License", "Programming Language :: Python :: 3.7", "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", "Topic :: Scientific/Engineering", "Topic :: Scientific/Engineering :: Artificial Intelligence", "Topic :: Software Development", diff --git a/plugins/flytekit-hive/setup.py b/plugins/flytekit-hive/setup.py index dcd5181d76..c3acd9e502 100644 --- a/plugins/flytekit-hive/setup.py +++ b/plugins/flytekit-hive/setup.py @@ -25,6 +25,8 @@ "License :: OSI Approved :: Apache Software License", "Programming Language :: Python :: 3.7", "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", "Topic :: Scientific/Engineering", "Topic :: Scientific/Engineering :: Artificial Intelligence", "Topic :: Software Development", diff --git a/plugins/flytekit-k8s-pod/setup.py b/plugins/flytekit-k8s-pod/setup.py index 2165a8034b..3b46d53b1f 100644 --- a/plugins/flytekit-k8s-pod/setup.py +++ b/plugins/flytekit-k8s-pod/setup.py @@ -28,6 +28,8 @@ "License :: OSI Approved :: Apache Software License", "Programming Language :: Python :: 3.7", "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", "Topic :: Scientific/Engineering", "Topic :: Scientific/Engineering :: Artificial Intelligence", "Topic :: Software Development", diff --git a/plugins/flytekit-kf-mpi/setup.py b/plugins/flytekit-kf-mpi/setup.py index 012b4e498c..2bd28ce387 100644 --- a/plugins/flytekit-kf-mpi/setup.py +++ b/plugins/flytekit-kf-mpi/setup.py @@ -25,6 +25,8 @@ "License :: OSI Approved :: Apache Software License", "Programming Language :: Python :: 3.7", "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", "Topic :: Scientific/Engineering", "Topic :: Scientific/Engineering :: Artificial Intelligence", "Topic :: Software Development", diff --git a/plugins/flytekit-kf-pytorch/setup.py b/plugins/flytekit-kf-pytorch/setup.py index 1fced6cf25..27f6c88c65 100644 --- a/plugins/flytekit-kf-pytorch/setup.py +++ b/plugins/flytekit-kf-pytorch/setup.py @@ -25,6 +25,8 @@ "License :: OSI Approved :: Apache Software License", "Programming Language :: Python :: 3.7", "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", "Topic :: Scientific/Engineering", "Topic :: Scientific/Engineering :: Artificial Intelligence", "Topic :: Software Development", diff --git a/plugins/flytekit-kf-tensorflow/setup.py b/plugins/flytekit-kf-tensorflow/setup.py index f4c8cb7f82..ad479d93e8 100644 --- a/plugins/flytekit-kf-tensorflow/setup.py +++ b/plugins/flytekit-kf-tensorflow/setup.py @@ -26,6 +26,8 @@ "License :: OSI Approved :: Apache Software License", "Programming Language :: Python :: 3.7", "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", "Topic :: Scientific/Engineering", "Topic :: Scientific/Engineering :: Artificial Intelligence", "Topic :: Software Development", diff --git a/plugins/flytekit-pandera/setup.py b/plugins/flytekit-pandera/setup.py index efc13ecc7f..fadbbea01a 100644 --- a/plugins/flytekit-pandera/setup.py +++ b/plugins/flytekit-pandera/setup.py @@ -25,6 +25,8 @@ "License :: OSI Approved :: Apache Software License", "Programming Language :: Python :: 3.7", "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", "Topic :: Scientific/Engineering", "Topic :: Scientific/Engineering :: Artificial Intelligence", "Topic :: Software Development", diff --git a/plugins/flytekit-papermill/setup.py b/plugins/flytekit-papermill/setup.py index 7c1c295aed..4fa4943302 100644 --- a/plugins/flytekit-papermill/setup.py +++ b/plugins/flytekit-papermill/setup.py @@ -31,6 +31,8 @@ "License :: OSI Approved :: Apache Software License", "Programming Language :: Python :: 3.7", "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", "Topic :: Scientific/Engineering", "Topic :: Scientific/Engineering :: Artificial Intelligence", "Topic :: Software Development", diff --git a/plugins/flytekit-snowflake/setup.py b/plugins/flytekit-snowflake/setup.py index d1f79cd06e..455444f84f 100644 --- a/plugins/flytekit-snowflake/setup.py +++ b/plugins/flytekit-snowflake/setup.py @@ -25,6 +25,8 @@ "License :: OSI Approved :: Apache Software License", "Programming Language :: Python :: 3.7", "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", "Topic :: Scientific/Engineering", "Topic :: Scientific/Engineering :: Artificial Intelligence", "Topic :: Software Development", diff --git a/plugins/flytekit-spark/setup.py b/plugins/flytekit-spark/setup.py index 655d5bbfbc..ca98542a16 100644 --- a/plugins/flytekit-spark/setup.py +++ b/plugins/flytekit-spark/setup.py @@ -25,6 +25,8 @@ "License :: OSI Approved :: Apache Software License", "Programming Language :: Python :: 3.7", "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", "Topic :: Scientific/Engineering", "Topic :: Scientific/Engineering :: Artificial Intelligence", "Topic :: Software Development", diff --git a/plugins/flytekit-sqlalchemy/setup.py b/plugins/flytekit-sqlalchemy/setup.py index e39b139552..1c5d1524db 100644 --- a/plugins/flytekit-sqlalchemy/setup.py +++ b/plugins/flytekit-sqlalchemy/setup.py @@ -25,6 +25,8 @@ "License :: OSI Approved :: Apache Software License", "Programming Language :: Python :: 3.7", "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", "Topic :: Scientific/Engineering", "Topic :: Scientific/Engineering :: Artificial Intelligence", "Topic :: Software Development", diff --git a/setup.py b/setup.py index f781690027..5bf9f83aa7 100644 --- a/setup.py +++ b/setup.py @@ -112,6 +112,7 @@ "Programming Language :: Python :: 3.7", "Programming Language :: Python :: 3.8", "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", "Topic :: Scientific/Engineering", "Topic :: Scientific/Engineering :: Artificial Intelligence", "Topic :: Software Development", From 2aa37625cac19eb4b83c7b52cef81b76c694e750 Mon Sep 17 00:00:00 2001 From: Samhita Alla Date: Tue, 4 Jan 2022 10:01:53 +0530 Subject: [PATCH 047/128] add `with_overrides` to map task (#794) * add with_overrides Signed-off-by: Samhita Alla * remove Resources Signed-off-by: Samhita Alla Signed-off-by: maximsmol --- flytekit/core/map_task.py | 7 ++++--- tests/flytekit/unit/core/test_map_task.py | 12 ++++-------- 2 files changed, 8 insertions(+), 11 deletions(-) diff --git a/flytekit/core/map_task.py b/flytekit/core/map_task.py index 60061cbb0e..42645a4797 100644 --- a/flytekit/core/map_task.py +++ b/flytekit/core/map_task.py @@ -207,7 +207,7 @@ def _raw_execute(self, **kwargs) -> Any: def map_task(task_function: PythonFunctionTask, concurrency: int = None, min_success_ratio: float = None, **kwargs): """ - Use a map task for parallelizable tasks that are run across a List of an input type. A map task can be composed of + Use a map task for parallelizable tasks that run across a list of an input type. A map task can be composed of any individual :py:class:`flytekit.PythonFunctionTask`. Invoke a map task with arguments using the :py:class:`list` version of the expected input. @@ -220,8 +220,8 @@ def map_task(task_function: PythonFunctionTask, concurrency: int = None, min_suc :language: python :dedent: 4 - At run time, the underlying map task will be run for every value in the input collection. Task-specific attributes - such as :py:class:`flytekit.TaskMetadata` and :py:class:`flytekit.Resources` are applied to individual instances + At run time, the underlying map task will be run for every value in the input collection. Attributes + such as :py:class:`flytekit.TaskMetadata` and ``with_overrides`` are applied to individual instances of the mapped task. :param task_function: This argument is implicitly passed and represents the repeatable function @@ -230,6 +230,7 @@ def map_task(task_function: PythonFunctionTask, concurrency: int = None, min_suc all inputs are processed. :param min_success_ratio: If specified, this determines the minimum fraction of total jobs which can complete successfully before terminating this task and marking it successful. + """ if not isinstance(task_function, PythonFunctionTask): raise ValueError( diff --git a/tests/flytekit/unit/core/test_map_task.py b/tests/flytekit/unit/core/test_map_task.py index 95df669829..31cbceffe9 100644 --- a/tests/flytekit/unit/core/test_map_task.py +++ b/tests/flytekit/unit/core/test_map_task.py @@ -3,7 +3,7 @@ import pytest -from flytekit import LaunchPlan, Resources, map_task +from flytekit import LaunchPlan, map_task from flytekit.common.translator import get_serializable from flytekit.core import context_manager from flytekit.core.context_manager import Image, ImageConfig @@ -33,13 +33,9 @@ def my_mappable_task(a: int) -> str: @workflow def my_wf(x: typing.List[int]) -> typing.List[str]: - return map_task( - my_mappable_task, - metadata=TaskMetadata(retries=1), - requests=Resources(cpu="10M"), - concurrency=10, - min_success_ratio=0.75, - )(a=x) + return map_task(my_mappable_task, metadata=TaskMetadata(retries=1), concurrency=10, min_success_ratio=0.75,)( + a=x + ).with_overrides(cpu="10M") # test_map_task_end From 2109643d0868a6f10890e56f8f51b536bf8c406c Mon Sep 17 00:00:00 2001 From: Samhita Alla Date: Fri, 7 Jan 2022 23:33:18 +0530 Subject: [PATCH 048/128] bump docsearch version (#805) Signed-off-by: Samhita Alla Signed-off-by: maximsmol --- doc-requirements.txt | 57 ++++++++++++++++++++++---------------------- 1 file changed, 29 insertions(+), 28 deletions(-) diff --git a/doc-requirements.txt b/doc-requirements.txt index cba57b13b7..6ad35bfea2 100644 --- a/doc-requirements.txt +++ b/doc-requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.9 +# This file is autogenerated by pip-compile with python 3.8 # To update, run: # # make doc-requirements.txt @@ -10,11 +10,15 @@ alabaster==0.7.12 # via sphinx ansiwrap==0.8.4 # via papermill +appnope==0.1.2 + # via + # ipykernel + # ipython arrow==1.2.1 # via jinja2-time -astroid==2.9.0 +astroid==2.9.2 # via sphinx-autoapi -attrs==21.2.0 +attrs==21.4.0 # via jsonschema babel==2.9.1 # via sphinx @@ -33,9 +37,9 @@ black==21.12b0 # via papermill bleach==4.1.0 # via nbconvert -boto3==1.20.26 +boto3==1.20.30 # via sagemaker-training -botocore==1.23.26 +botocore==1.23.30 # via # boto3 # s3transfer @@ -48,7 +52,7 @@ cffi==1.15.0 # pynacl chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.9 +charset-normalizer==2.0.10 # via requests checksumdir==1.2.0 # via flytekit @@ -69,12 +73,11 @@ cryptography==36.0.1 # via # -r doc-requirements.in # paramiko - # secretstorage css-html-js-minify==2.5.5 # via sphinx-material dataclasses-json==0.5.6 # via flytekit -decorator==5.1.0 +decorator==5.1.1 # via # ipython # retry @@ -82,7 +85,7 @@ defusedxml==0.7.1 # via nbconvert deprecated==1.2.13 # via flytekit -diskcache==5.3.0 +diskcache==5.4.0 # via flytekit docker-image-py==0.1.12 # via flytekit @@ -95,7 +98,7 @@ entrypoints==0.3 # jupyter-client # nbconvert # papermill -flyteidl==0.21.13 +flyteidl==0.21.17 # via flytekit furo @ git+git://github.com/flyteorg/furo@main # via -r doc-requirements.in @@ -115,11 +118,13 @@ imagesize==1.3.0 # via sphinx importlib-metadata==4.10.0 # via keyring +importlib-resources==5.4.0 + # via jsonschema inotify_simple==1.2.1 # via sagemaker-training ipykernel==5.5.6 # via flytekit -ipython==7.30.1 +ipython==7.31.0 # via ipykernel ipython-genutils==0.2.0 # via @@ -127,10 +132,6 @@ ipython-genutils==0.2.0 # nbformat jedi==0.18.1 # via ipython -jeepney==0.7.1 - # via - # keyring - # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -144,7 +145,7 @@ jmespath==0.10.0 # via # boto3 # botocore -jsonschema==4.3.2 +jsonschema==4.3.3 # via nbformat jupyter-client==7.1.0 # via @@ -159,7 +160,7 @@ jupyterlab-pygments==0.1.2 # via nbconvert k8s-proto==0.0.3 # via flytekit -keyring==23.4.0 +keyring==23.5.0 # via flytekit lazy-object-proxy==1.7.1 # via astroid @@ -190,7 +191,7 @@ nbclient==0.5.9 # via # nbconvert # papermill -nbconvert==6.3.0 +nbconvert==6.4.0 # via flytekit nbformat==5.1.3 # via @@ -201,7 +202,7 @@ nest-asyncio==1.5.4 # via # jupyter-client # nbclient -numpy==1.21.5 +numpy==1.22.0 # via # flytekit # pandas @@ -218,7 +219,7 @@ pandocfilters==1.5.0 # via nbconvert papermill==2.3.3 # via flytekit -paramiko==2.8.1 +paramiko==2.9.1 # via sagemaker-training parso==0.8.3 # via jedi @@ -228,7 +229,7 @@ pexpect==4.8.0 # via ipython pickleshare==0.7.5 # via ipython -platformdirs==2.4.0 +platformdirs==2.4.1 # via black poyo==0.5.0 # via cookiecutter @@ -240,7 +241,7 @@ protobuf==3.19.1 # flytekit # k8s-proto # sagemaker-training -psutil==5.8.0 +psutil==5.9.0 # via sagemaker-training ptyprocess==0.7.0 # via pexpect @@ -252,7 +253,7 @@ pyarrow==6.0.1 # via flytekit pycparser==2.21 # via cffi -pygments==2.10.0 +pygments==2.11.2 # via # ipython # jupyterlab-pygments @@ -296,7 +297,7 @@ pyzmq==22.3.0 # via jupyter-client regex==2021.11.10 # via docker-image-py -requests==2.26.0 +requests==2.27.1 # via # cookiecutter # flytekit @@ -315,8 +316,6 @@ sagemaker-training==3.9.2 # via flytekit scipy==1.7.3 # via sagemaker-training -secretstorage==3.3.1 - # via keyring six==1.16.0 # via # bcrypt @@ -437,8 +436,10 @@ wrapt==1.13.3 # astroid # deprecated # flytekit -zipp==3.6.0 - # via importlib-metadata +zipp==3.7.0 + # via + # importlib-metadata + # importlib-resources zope.event==4.5.0 # via gevent zope.interface==5.4.0 From ec85174611c5bb2d1d13cd414c31ce293402c74a Mon Sep 17 00:00:00 2001 From: Niels Bantilan Date: Fri, 7 Jan 2022 19:18:50 -0500 Subject: [PATCH 049/128] update docs for new navbar theme (#806) Signed-off-by: maximsmol --- docs/source/conf.py | 5 +++-- docs/source/index.rst | 2 ++ 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/docs/source/conf.py b/docs/source/conf.py index 5d0c11a387..00d6803fcf 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -20,7 +20,7 @@ # -- Project information ----------------------------------------------------- -project = "Flyte Cookbook [Python]" +project = "Flytekit" copyright = "2021, Flyte" author = "Flyte" @@ -90,7 +90,7 @@ # a list of builtin themes. # html_theme = "furo" -html_title = "Flyte Docs" +html_title = "Flyte" html_theme_options = { "light_css_variables": { @@ -126,6 +126,7 @@ # to template names. # html_logo = "flyte_circle_gradient_1_4x4.png" +html_favicon = "flyte_circle_gradient_1_4x4.png" pygments_style = "tango" pygments_dark_style = "native" diff --git a/docs/source/index.rst b/docs/source/index.rst index 2fee11c40b..f83cbaa9b2 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -62,6 +62,8 @@ Expected output: |book| API Reference |hands-helping| Community +.. NOTE: the caption text is important for the sphinx theme to correctly render the nav header +.. https://github.com/flyteorg/furo .. toctree:: :maxdepth: -1 :caption: Flytekit SDK From 279ce1db7831172c009c9eff6a0b37e9406d48f3 Mon Sep 17 00:00:00 2001 From: Yuvraj Date: Wed, 12 Jan 2022 00:52:19 +0530 Subject: [PATCH 050/128] fix requirment.txt github issue (#810) Signed-off-by: Yuvraj Signed-off-by: maximsmol --- dev-requirements.in | 2 +- dev-requirements.txt | 2 +- doc-requirements.in | 2 +- doc-requirements.txt | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/dev-requirements.in b/dev-requirements.in index 9743d0fc20..c8eb2c0601 100644 --- a/dev-requirements.in +++ b/dev-requirements.in @@ -1,6 +1,6 @@ -c requirements.txt -git+git://github.com/flyteorg/pytest-flyte@main#egg=pytest-flyte +git+https://github.com/flyteorg/pytest-flyte@main#egg=pytest-flyte coverage[toml] joblib mock diff --git a/dev-requirements.txt b/dev-requirements.txt index 8cc4f7ba3f..294d19a69f 100644 --- a/dev-requirements.txt +++ b/dev-requirements.txt @@ -262,7 +262,7 @@ pytest==6.2.5 # pytest-flyte pytest-docker==0.10.3 # via pytest-flyte -pytest-flyte @ git+git://github.com/flyteorg/pytest-flyte@main +pytest-flyte @ git+https://github.com/flyteorg/pytest-flyte@main # via -r dev-requirements.in python-dateutil==2.8.1 # via diff --git a/doc-requirements.in b/doc-requirements.in index 9cc60cfd43..972e5cf166 100644 --- a/doc-requirements.in +++ b/doc-requirements.in @@ -1,7 +1,7 @@ .[all] -e file:.#egg=flytekit -git+git://github.com/flyteorg/furo@main +git+https://github.com/flyteorg/furo@main sphinx sphinx-gallery sphinx-prompt diff --git a/doc-requirements.txt b/doc-requirements.txt index 6ad35bfea2..265a311b28 100644 --- a/doc-requirements.txt +++ b/doc-requirements.txt @@ -100,7 +100,7 @@ entrypoints==0.3 # papermill flyteidl==0.21.17 # via flytekit -furo @ git+git://github.com/flyteorg/furo@main +furo @ git+https://github.com/flyteorg/furo@main # via -r doc-requirements.in gevent==21.12.0 # via sagemaker-training From fa50c9a8fc16fe7805727dc0287716accaac6fa2 Mon Sep 17 00:00:00 2001 From: Niels Bantilan Date: Wed, 12 Jan 2022 18:50:58 -0500 Subject: [PATCH 051/128] Add sphinx panels (#815) Signed-off-by: Ketan Umare Signed-off-by: maximsmol --- doc-requirements.in | 1 + doc-requirements.txt | 43 +++++++++++++++++++++++++++++-------------- docs/source/conf.py | 1 + 3 files changed, 31 insertions(+), 14 deletions(-) diff --git a/doc-requirements.in b/doc-requirements.in index 972e5cf166..4be8ade6d6 100644 --- a/doc-requirements.in +++ b/doc-requirements.in @@ -10,6 +10,7 @@ sphinx-code-include sphinx-autoapi sphinx-copybutton sphinx_fontawesome +sphinx-panels sphinxcontrib-yt grpcio cryptography diff --git a/doc-requirements.txt b/doc-requirements.txt index 265a311b28..9590c1face 100644 --- a/doc-requirements.txt +++ b/doc-requirements.txt @@ -16,8 +16,10 @@ appnope==0.1.2 # ipython arrow==1.2.1 # via jinja2-time -astroid==2.9.2 +astroid==2.9.3 # via sphinx-autoapi +asttokens==2.0.5 + # via stack-data attrs==21.4.0 # via jsonschema babel==2.9.1 @@ -34,12 +36,14 @@ beautifulsoup4==4.10.0 binaryornot==0.4.4 # via cookiecutter black==21.12b0 - # via papermill + # via + # ipython + # papermill bleach==4.1.0 # via nbconvert -boto3==1.20.30 +boto3==1.20.34 # via sagemaker-training -botocore==1.23.30 +botocore==1.23.34 # via # boto3 # s3transfer @@ -92,13 +96,17 @@ docker-image-py==0.1.12 docstring-parser==0.13 # via flytekit docutils==0.17.1 - # via sphinx + # via + # sphinx + # sphinx-panels entrypoints==0.3 # via # jupyter-client # nbconvert # papermill -flyteidl==0.21.17 +executing==0.8.2 + # via stack-data +flyteidl==0.21.22 # via flytekit furo @ git+https://github.com/flyteorg/furo@main # via -r doc-requirements.in @@ -124,7 +132,7 @@ inotify_simple==1.2.1 # via sagemaker-training ipykernel==5.5.6 # via flytekit -ipython==7.31.0 +ipython==8.0.0 # via ipykernel ipython-genutils==0.2.0 # via @@ -145,7 +153,7 @@ jmespath==0.10.0 # via # boto3 # botocore -jsonschema==4.3.3 +jsonschema==4.4.0 # via nbformat jupyter-client==7.1.0 # via @@ -219,7 +227,7 @@ pandocfilters==1.5.0 # via nbconvert papermill==2.3.3 # via flytekit -paramiko==2.9.1 +paramiko==2.9.2 # via sagemaker-training parso==0.8.3 # via jedi @@ -235,7 +243,7 @@ poyo==0.5.0 # via cookiecutter prompt-toolkit==3.0.24 # via ipython -protobuf==3.19.1 +protobuf==3.19.3 # via # flyteidl # flytekit @@ -245,6 +253,8 @@ psutil==5.9.0 # via sagemaker-training ptyprocess==0.7.0 # via pexpect +pure-eval==0.2.1 + # via stack-data py==1.11.0 # via retry py4j==0.10.9.2 @@ -260,7 +270,7 @@ pygments==2.11.2 # nbconvert # sphinx # sphinx-prompt -pynacl==1.4.0 +pynacl==1.5.0 # via paramiko pyparsing==3.0.6 # via packaging @@ -304,7 +314,7 @@ requests==2.27.1 # papermill # responses # sphinx -responses==0.16.0 +responses==0.17.0 # via flytekit retry==0.9.2 # via flytekit @@ -318,12 +328,12 @@ scipy==1.7.3 # via sagemaker-training six==1.16.0 # via + # asttokens # bcrypt # bleach # cookiecutter # flytekit # grpcio - # pynacl # python-dateutil # responses # retrying @@ -346,6 +356,7 @@ sphinx==4.3.2 # sphinx-fontawesome # sphinx-gallery # sphinx-material + # sphinx-panels # sphinx-prompt # sphinxcontrib-yt sphinx-autoapi==1.8.4 @@ -360,6 +371,8 @@ sphinx-gallery==0.10.1 # via -r doc-requirements.in sphinx-material==0.0.35 # via -r doc-requirements.in +sphinx-panels==0.6.0 + # via -r doc-requirements.in sphinx-prompt==1.5.0 # via -r doc-requirements.in sphinxcontrib-applehelp==1.0.2 @@ -376,6 +389,8 @@ sphinxcontrib-serializinghtml==1.1.5 # via sphinx sphinxcontrib-yt==0.2.2 # via -r doc-requirements.in +stack-data==0.1.3 + # via ipython statsd==3.3.0 # via flytekit tenacity==8.0.1 @@ -417,7 +432,7 @@ unidecode==1.3.2 # via # python-slugify # sphinx-autoapi -urllib3==1.26.7 +urllib3==1.26.8 # via # botocore # flytekit diff --git a/docs/source/conf.py b/docs/source/conf.py index 00d6803fcf..9ca0b626be 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -47,6 +47,7 @@ "sphinx-prompt", "sphinx_copybutton", "sphinx_fontawesome", + "sphinx_panels", "sphinxcontrib.yt", ] From 71c922ae00f77b65456d5bdf690eea27717a3bd0 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Fri, 14 Jan 2022 02:18:15 +0800 Subject: [PATCH 052/128] Schema overhaul (#785) Signed-off-by: Kevin Su Signed-off-by: Yee Hing Tong --- dev-requirements.txt | 12 +- flytekit/__init__.py | 1 + flytekit/core/data_persistence.py | 2 +- flytekit/core/interface.py | 6 +- flytekit/core/type_engine.py | 41 +- flytekit/models/core/types.py | 2 + flytekit/models/literals.py | 78 ++- flytekit/models/types.py | 93 +++ flytekit/types/schema/types.py | 13 +- flytekit/types/structured/__init__.py | 21 + flytekit/types/structured/basic_dfs.py | 104 ++++ flytekit/types/structured/bigquery.py | 112 ++++ .../types/structured/structured_dataset.py | 584 ++++++++++++++++++ .../great_expectations/schema.py | 24 +- .../flytekitplugins/spark/__init__.py | 2 +- .../flytekitplugins/spark/schema.py | 47 +- plugins/flytekit-spark/tests/test_wf.py | 58 +- requirements-spark2.txt | 2 +- requirements.txt | 2 +- setup.py | 2 +- .../workflows/requirements.txt | 2 +- tests/flytekit/unit/core/test_flyte_file.py | 3 +- tests/flytekit/unit/core/test_imperative.py | 6 +- tests/flytekit/unit/core/test_interface.py | 15 + tests/flytekit/unit/core/test_local_cache.py | 12 + .../unit/core/test_structured_dataset.py | 253 ++++++++ .../core/test_structured_dataset_handlers.py | 43 ++ tests/flytekit/unit/core/test_type_delayed.py | 21 + tests/flytekit/unit/core/test_type_engine.py | 55 ++ tests/flytekit/unit/core/test_type_hints.py | 33 + tests/flytekit/unit/core/test_workflows.py | 63 ++ tests/flytekit/unit/models/test_literals.py | 37 ++ .../test_structured_dataset_workflow.py | 224 +++++++ 33 files changed, 1937 insertions(+), 36 deletions(-) create mode 100644 flytekit/types/structured/__init__.py create mode 100644 flytekit/types/structured/basic_dfs.py create mode 100644 flytekit/types/structured/bigquery.py create mode 100644 flytekit/types/structured/structured_dataset.py create mode 100644 tests/flytekit/unit/core/test_structured_dataset.py create mode 100644 tests/flytekit/unit/core/test_structured_dataset_handlers.py create mode 100644 tests/flytekit/unit/type_engines/structured_dataset/test_structured_dataset_workflow.py diff --git a/dev-requirements.txt b/dev-requirements.txt index 294d19a69f..4d5c69c79c 100644 --- a/dev-requirements.txt +++ b/dev-requirements.txt @@ -24,6 +24,7 @@ bcrypt==3.2.0 # via # -c requirements.txt # paramiko + # secretstorage binaryornot==0.4.4 # via # -c requirements.txt @@ -77,7 +78,6 @@ cryptography==36.0.1 # via # -c requirements.txt # paramiko - # secretstorage dataclasses-json==0.5.6 # via # -c requirements.txt @@ -118,7 +118,7 @@ docstring-parser==0.13 # flytekit filelock==3.4.0 # via virtualenv -flyteidl==0.21.13 +flyteidl==0.21.17 # via # -c requirements.txt # flytekit @@ -181,6 +181,10 @@ marshmallow-jsonschema==0.13.0 # via # -c requirements.txt # flytekit +secretstorage==3.3.1 + # via + # -c requirements.txt + # keyring mock==4.0.3 # via -r dev-requirements.in mypy==0.930 @@ -315,10 +319,6 @@ retry==0.9.2 # via # -c requirements.txt # flytekit -secretstorage==3.3.1 - # via - # -c requirements.txt - # keyring six==1.16.0 # via # -c requirements.txt diff --git a/flytekit/__init__.py b/flytekit/__init__.py index e1d51f6e4a..3a165084b0 100644 --- a/flytekit/__init__.py +++ b/flytekit/__init__.py @@ -187,6 +187,7 @@ from flytekit.models.literals import Blob, BlobMetadata, Literal, Scalar from flytekit.models.types import LiteralType from flytekit.types import directory, file, schema +from flytekit.types.structured.structured_dataset import StructuredDataset, StructuredDatasetType __version__ = "0.0.0+develop" diff --git a/flytekit/core/data_persistence.py b/flytekit/core/data_persistence.py index ad62cd0d18..b121d053fa 100644 --- a/flytekit/core/data_persistence.py +++ b/flytekit/core/data_persistence.py @@ -141,7 +141,7 @@ def find_plugin(cls, path: str) -> typing.Type[DataPersistence]: Returns a plugin for the given protocol, else raise a TypeError """ for k, p in cls._PLUGINS.items(): - if path.startswith(k): + if path.startswith(k) or path.startswith(k.replace("://", "")): return p raise TypeError(f"No plugin found for matching protocol of path {path}") diff --git a/flytekit/core/interface.py b/flytekit/core/interface.py index 216e337282..263ac05fb7 100644 --- a/flytekit/core/interface.py +++ b/flytekit/core/interface.py @@ -275,7 +275,11 @@ def transform_function_to_interface(fn: Callable, docstring: Optional[Docstring] For now the fancy object, maybe in the future a dumb object. """ - type_hints = typing.get_type_hints(fn) + try: + # include_extras can only be used in python >= 3.9 + type_hints = typing.get_type_hints(fn, include_extras=True) + except TypeError: + type_hints = typing.get_type_hints(fn) signature = inspect.signature(fn) return_annotation = type_hints.get("return", None) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 21c0749fbb..9c38120e69 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -11,6 +11,11 @@ from re import L from typing import NamedTuple, Optional, Type, cast +try: + from typing import Annotated, get_args, get_origin +except ImportError: + from typing_extensions import Annotated, get_origin, get_args + from dataclasses_json import DataClassJsonMixin, dataclass_json from google.protobuf import json_format as _json_format from google.protobuf import reflection as _proto_reflection @@ -38,10 +43,11 @@ Primitive, Scalar, Schema, + StructuredDatasetMetadata, Union, Void, ) -from flytekit.models.types import LiteralType, SimpleType, TypeStructure, UnionType +from flytekit.models.types import LiteralType, SimpleType, StructuredDatasetType, TypeStructure, UnionType try: from typing import get_args as _get_args @@ -308,6 +314,7 @@ def _serialize_flyte_type(self, python_val: T, python_type: Type[T]): from flytekit.types.directory.types import FlyteDirectory from flytekit.types.file import FlyteFile from flytekit.types.schema.types import FlyteSchema + from flytekit.types.structured.structured_dataset import StructuredDataset for f in dataclasses.fields(python_type): v = python_val.__getattribute__(f.name) @@ -316,6 +323,7 @@ def _serialize_flyte_type(self, python_val: T, python_type: Type[T]): issubclass(field_type, FlyteSchema) or issubclass(field_type, FlyteFile) or issubclass(field_type, FlyteDirectory) + or issubclass(field_type, StructuredDataset) ): lv = TypeEngine.to_literal(FlyteContext.current_context(), v, field_type, None) # dataclass_json package will extract the "path" from FlyteFile, FlyteDirectory, and write it to a @@ -328,6 +336,13 @@ def _serialize_flyte_type(self, python_val: T, python_type: Type[T]): # as determined by the transformer. if issubclass(field_type, FlyteFile) or issubclass(field_type, FlyteDirectory): python_val.__setattr__(f.name, field_type(path=lv.scalar.blob.uri)) + elif issubclass(field_type, StructuredDataset): + python_val.__setattr__( + f.name, + field_type( + uri=lv.scalar.structured_dataset.uri, + ), + ) elif dataclasses.is_dataclass(field_type): self._serialize_flyte_type(v, field_type) @@ -336,6 +351,7 @@ def _deserialize_flyte_type(self, python_val: T, expected_python_type: Type) -> from flytekit.types.directory.types import FlyteDirectory, FlyteDirToMultipartBlobTransformer from flytekit.types.file.file import FlyteFile, FlyteFilePathTransformer from flytekit.types.schema.types import FlyteSchema, FlyteSchemaTransformer + from flytekit.types.structured.structured_dataset import StructuredDataset, StructuredDatasetTransformerEngine if not dataclasses.is_dataclass(expected_python_type): return python_val @@ -381,6 +397,21 @@ def _deserialize_flyte_type(self, python_val: T, expected_python_type: Type) -> ), expected_python_type, ) + elif issubclass(expected_python_type, StructuredDataset): + return StructuredDatasetTransformerEngine().to_python_value( + FlyteContext.current_context(), + Literal( + scalar=Scalar( + structured_dataset=StructuredDataset( + metadata=StructuredDatasetMetadata( + structured_dataset_type=StructuredDatasetType(format=python_val.file_format) + ), + uri=python_val.uri, + ) + ) + ), + expected_python_type, + ) else: for f in dataclasses.fields(expected_python_type): value = python_val.__getattribute__(f.name) @@ -528,6 +559,11 @@ def register_restricted_type( cls._RESTRICTED_TYPES.append(type) cls.register(RestrictedTypeTransformer(name, type)) + @classmethod + def register_additional_type(cls, transformer: TypeTransformer, additional_type: Type, override=False): + if additional_type not in cls._REGISTRY or override: + cls._REGISTRY[additional_type] = transformer + @classmethod def get_transformer(cls, python_type: Type) -> TypeTransformer[T]: """ @@ -553,6 +589,9 @@ def get_transformer(cls, python_type: Type) -> TypeTransformer[T]: """ # Step 1 + if get_origin(python_type) is Annotated: + python_type = get_args(python_type)[0] + if python_type in cls._REGISTRY: return cls._REGISTRY[python_type] diff --git a/flytekit/models/core/types.py b/flytekit/models/core/types.py index 784b789d90..4508961bbc 100644 --- a/flytekit/models/core/types.py +++ b/flytekit/models/core/types.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import typing from flyteidl.core import types_pb2 as _types_pb2 diff --git a/flytekit/models/literals.py b/flytekit/models/literals.py index b3acc7abde..9333f38ba2 100644 --- a/flytekit/models/literals.py +++ b/flytekit/models/literals.py @@ -1,3 +1,4 @@ +import typing from datetime import datetime as _datetime import pytz as _pytz @@ -10,6 +11,7 @@ from flytekit.models.types import LiteralType as _LiteralType from flytekit.models.types import OutputReference as _OutputReference from flytekit.models.types import SchemaType as _SchemaType +from flytekit.models.types import StructuredDatasetType class RetryStrategy(_common.FlyteIdlEntity): @@ -584,7 +586,58 @@ def from_flyte_idl(cls, pb2_object): :param flyteidl.core.literals_pb2.Schema pb2_object: :rtype: Schema """ - return cls(value=Literal.from_flyte_idl(pb2_object.value), stored_type=_LiteralType.from_flyte_idl(pb2_object.type)) + return cls( + value=Literal.from_flyte_idl(pb2_object.value), stored_type=_LiteralType.from_flyte_idl(pb2_object.type) + ) + + +class StructuredDatasetMetadata(_common.FlyteIdlEntity): + def __init__(self, structured_dataset_type: StructuredDatasetType = None): + self._structured_dataset_type = structured_dataset_type + + @property + def structured_dataset_type(self) -> StructuredDatasetType: + return self._structured_dataset_type + + def to_flyte_idl(self) -> _literals_pb2.StructuredDatasetMetadata: + return _literals_pb2.StructuredDatasetMetadata( + structured_dataset_type=self.structured_dataset_type.to_flyte_idl() + if self._structured_dataset_type + else None, + ) + + @classmethod + def from_flyte_idl(cls, pb2_object: _literals_pb2.StructuredDatasetMetadata) -> "StructuredDatasetMetadata": + return cls( + structured_dataset_type=StructuredDatasetType.from_flyte_idl(pb2_object.structured_dataset_type), + ) + + +class StructuredDataset(_common.FlyteIdlEntity): + def __init__(self, uri: str, metadata: typing.Optional[StructuredDatasetMetadata] = None): + """ + A strongly typed schema that defines the interface of data retrieved from the underlying storage medium. + """ + self._uri = uri + self._metadata = metadata + + @property + def uri(self) -> str: + return self._uri + + @property + def metadata(self) -> StructuredDatasetMetadata: + return self._metadata + + def to_flyte_idl(self) -> _literals_pb2.StructuredDataset: + return _literals_pb2.StructuredDataset( + uri=self.uri, metadata=self.metadata.to_flyte_idl() if self.metadata else None + ) + + @classmethod + def from_flyte_idl(cls, pb2_object: _literals_pb2.StructuredDataset) -> "StructuredDataset": + return cls(uri=pb2_object.uri, metadata=StructuredDatasetMetadata.from_flyte_idl(pb2_object.metadata)) + class LiteralCollection(_common.FlyteIdlEntity): def __init__(self, literals): @@ -656,6 +709,7 @@ def __init__( none_type: Void = None, error=None, generic: Struct = None, + structured_dataset: StructuredDataset = None, ): """ Scalar wrapper around Flyte types. Only one can be specified. @@ -667,6 +721,7 @@ def __init__( :param Void none_type: :param error: :param google.protobuf.struct_pb2.Struct generic: + :param StructuredDataset structured_dataset: """ self._primitive = primitive @@ -677,6 +732,7 @@ def __init__( self._none_type = none_type self._error = error self._generic = generic + self._structured_dataset = structured_dataset @property def primitive(self): @@ -734,13 +790,27 @@ def generic(self): """ return self._generic + @property + def structured_dataset(self) -> StructuredDataset: + return self._structured_dataset + @property def value(self): """ Returns whichever value is set :rtype: T """ - return self.primitive or self.blob or self.binary or self.schema or self.union or self.none_type or self.error + return ( + self.primitive + or self.blob + or self.binary + or self.schema + or self.union + or self.none_type + or self.error + or self.generic + or self.structured_dataset + ) def to_flyte_idl(self): """ @@ -755,6 +825,7 @@ def to_flyte_idl(self): none_type=self.none_type.to_flyte_idl() if self.none_type is not None else None, error=self.error if self.error is not None else None, generic=self.generic, + structured_dataset=self.structured_dataset.to_flyte_idl() if self.structured_dataset is not None else None, ) @classmethod @@ -773,6 +844,9 @@ def from_flyte_idl(cls, pb2_object): none_type=Void.from_flyte_idl(pb2_object.none_type) if pb2_object.HasField("none_type") else None, error=pb2_object.error if pb2_object.HasField("error") else None, generic=pb2_object.generic if pb2_object.HasField("generic") else None, + structured_dataset=StructuredDataset.from_flyte_idl(pb2_object.structured_dataset) + if pb2_object.HasField("structured_dataset") + else None, ) diff --git a/flytekit/models/types.py b/flytekit/models/types.py index d967e1bfb4..a3d36de77b 100644 --- a/flytekit/models/types.py +++ b/flytekit/models/types.py @@ -143,6 +143,87 @@ def from_flyte_idl(cls, proto: _types_pb2.TypeStructure): return cls(tag=proto.tag) +class StructuredDatasetType(_common.FlyteIdlEntity): + class DatasetColumn(_common.FlyteIdlEntity): + def __init__(self, name: str, literal_type: "LiteralType"): + self._name = name + self._literal_type = literal_type + + @property + def name(self) -> str: + """ + Name for the column + """ + return self._name + + @property + def literal_type(self) -> "LiteralType": + """ + A LiteralType that defines the type of this column + """ + return self._literal_type + + def to_flyte_idl(self) -> _types_pb2.StructuredDatasetType.DatasetColumn: + return _types_pb2.StructuredDatasetType.DatasetColumn( + name=self.name, literal_type=self.literal_type.to_flyte_idl() + ) + + @classmethod + def from_flyte_idl( + cls, proto: _types_pb2.StructuredDatasetType.DatasetColumn + ) -> _types_pb2.StructuredDatasetType.DatasetColumn: + return cls(name=proto.name, literal_type=LiteralType.from_flyte_idl(proto.literal_type)) + + def __init__( + self, + columns: typing.List[DatasetColumn] = None, + format: str = "", + external_schema_type: str = None, + external_schema_bytes: bytes = None, + ): + self._columns = columns + self._format = format + self._external_schema_type = external_schema_type + self._external_schema_bytes = external_schema_bytes + + @property + def columns(self) -> typing.List[DatasetColumn]: + return self._columns + + @property + def format(self) -> str: + return self._format + + @format.setter + def format(self, format: str): + self._format = format + + @property + def external_schema_type(self) -> str: + return self._external_schema_type + + @property + def external_schema_bytes(self) -> bytes: + return self._external_schema_bytes + + def to_flyte_idl(self) -> _types_pb2.StructuredDatasetType: + return _types_pb2.StructuredDatasetType( + columns=[c.to_flyte_idl() for c in self.columns] if self.columns else None, + format=self.format, + external_schema_type=self.external_schema_type if self.external_schema_type else None, + external_schema_bytes=self.external_schema_bytes if self.external_schema_bytes else None, + ) + + @classmethod + def from_flyte_idl(cls, proto: _types_pb2.StructuredDatasetType) -> _types_pb2.StructuredDatasetType: + return cls( + columns=[StructuredDatasetType.DatasetColumn.from_flyte_idl(c) for c in proto.columns], + format=proto.format, + external_schema_type=proto.external_schema_type, + external_schema_bytes=proto.external_schema_bytes, + ) + + class LiteralType(_common.FlyteIdlEntity): def __init__( self, @@ -154,6 +235,7 @@ def __init__( enum_type=None, union_type=None, structure=None, + structured_dataset_type=None, metadata=None, ): """ @@ -168,6 +250,7 @@ def __init__( :param flytekit.models.core.types.EnumType enum_type: For enum objects, describes an enum :param flytekit.models.core.types.UnionType union_type: For union objects, describes an python union type. :param flytekit.models.core.types.TypeStructure structure: Type matching hints + :param flytekit.models.core.types.StructuredDatasetType structured_dataset_type: structured dataset :param dict[Text, T] metadata: Additional data describing the type """ self._simple = simple @@ -178,6 +261,7 @@ def __init__( self._enum_type = enum_type self._union_type = union_type self._structure = structure + self._structured_dataset_type = structured_dataset_type self._metadata = metadata @property @@ -218,6 +302,9 @@ def union_type(self) -> UnionType: def structure(self) -> TypeStructure: return self._structure + def structured_dataset_type(self) -> StructuredDatasetType: + return self._structured_dataset_type + @property def metadata(self): """ @@ -246,6 +333,9 @@ def to_flyte_idl(self): enum_type=self.enum_type.to_flyte_idl() if self.enum_type else None, union_type=self.union_type.to_flyte_idl() if self.union_type else None, structure=self.structure.to_flyte_idl() if self.structure else None, + structured_dataset_type=self.structured_dataset_type.to_flyte_idl() + if self.structured_dataset_type + else None, metadata=metadata, ) return t @@ -271,6 +361,9 @@ def from_flyte_idl(cls, proto): enum_type=_core_types.EnumType.from_flyte_idl(proto.enum_type) if proto.HasField("enum_type") else None, union_type=UnionType.from_flyte_idl(proto.union_type) if proto.HasField("union_type") else None, structure=TypeStructure.from_flyte_idl(proto.structure) if proto.HasField("structure") else None, + structured_dataset_type=StructuredDatasetType.from_flyte_idl(proto.structured_dataset_type) + if proto.HasField("structured_dataset_type") + else None, metadata=_json_format.MessageToDict(proto.metadata) or None, ) diff --git a/flytekit/types/schema/types.py b/flytekit/types/schema/types.py index 6b44af5396..b9d2a51372 100644 --- a/flytekit/types/schema/types.py +++ b/flytekit/types/schema/types.py @@ -379,12 +379,19 @@ def to_literal( return Literal(scalar=Scalar(schema=Schema(schema.remote_path, self._get_schema_type(python_type)))) def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[FlyteSchema]) -> FlyteSchema: - if not (lv and lv.scalar and lv.scalar.schema): - raise TypeTransformerFailedError("Can only convert a literal schema to a FlyteSchema") - def downloader(x, y): ctx.file_access.get_data(x, y, is_multipart=True) + if lv and lv.scalar and lv.scalar.structured_dataset: + return expected_python_type( + local_path=ctx.file_access.get_random_local_directory(), + remote_path=lv.scalar.structured_dataset.uri, + downloader=downloader, + supported_mode=SchemaOpenMode.READ, + ) + if not (lv and lv.scalar and lv.scalar.schema): + raise AssertionError("Can only convert a literal schema to a FlyteSchema") + return expected_python_type( local_path=ctx.file_access.get_random_local_directory(), remote_path=lv.scalar.schema.uri, diff --git a/flytekit/types/structured/__init__.py b/flytekit/types/structured/__init__.py new file mode 100644 index 0000000000..da80c89016 --- /dev/null +++ b/flytekit/types/structured/__init__.py @@ -0,0 +1,21 @@ +from flytekit.loggers import logger + +from .basic_dfs import ( + ArrowToParquetEncodingHandler, + PandasToParquetEncodingHandler, + ParquetToArrowDecodingHandler, + ParquetToPandasDecodingHandler, +) + +try: + from .bigquery import ( + ArrowToBQEncodingHandlers, + BQToArrowDecodingHandler, + BQToPandasDecodingHandler, + PandasToBQEncodingHandlers, + ) +except ImportError: + logger.info( + "We won't register bigquery handler for structured dataset because " + "we can't find the packages google-cloud-bigquery-storage and google-cloud-bigquery" + ) diff --git a/flytekit/types/structured/basic_dfs.py b/flytekit/types/structured/basic_dfs.py new file mode 100644 index 0000000000..f7a28bdc75 --- /dev/null +++ b/flytekit/types/structured/basic_dfs.py @@ -0,0 +1,104 @@ +import os +import typing +from typing import TypeVar + +import pandas +import pandas as pd +import pyarrow as pa +import pyarrow.parquet as pq + +from flytekit import FlyteContext +from flytekit.core.data_persistence import DataPersistencePlugins +from flytekit.models import literals +from flytekit.models.literals import StructuredDatasetMetadata +from flytekit.models.types import StructuredDatasetType +from flytekit.types.structured.structured_dataset import ( + FLYTE_DATASET_TRANSFORMER, + LOCAL, + PARQUET, + S3, + StructuredDataset, + StructuredDatasetDecoder, + StructuredDatasetEncoder, +) + +T = TypeVar("T") + + +class PandasToParquetEncodingHandler(StructuredDatasetEncoder): + def __init__(self, protocol: str): + super().__init__(pd.DataFrame, protocol, PARQUET) + # todo: Use this somehow instead of relaying ont he ctx file_access + self._persistence = DataPersistencePlugins.find_plugin(protocol)() + + def encode( + self, + ctx: FlyteContext, + structured_dataset: StructuredDataset, + structured_dataset_type: StructuredDatasetType, + ) -> literals.StructuredDataset: + + path = typing.cast(str, structured_dataset.uri) or ctx.file_access.get_random_remote_directory() + df = typing.cast(pd.DataFrame, structured_dataset.dataframe) + local_dir = ctx.file_access.get_random_local_directory() + local_path = os.path.join(local_dir, f"{0:05}") + df.to_parquet(local_path, coerce_timestamps="us", allow_truncated_timestamps=False) + ctx.file_access.upload_directory(local_dir, path) + structured_dataset_type.format = PARQUET + return literals.StructuredDataset(uri=path, metadata=StructuredDatasetMetadata(structured_dataset_type)) + + +class ParquetToPandasDecodingHandler(StructuredDatasetDecoder): + def __init__(self, protocol: str): + super().__init__(pd.DataFrame, protocol, PARQUET) + + def decode( + self, + ctx: FlyteContext, + flyte_value: literals.StructuredDataset, + ) -> pd.DataFrame: + path = flyte_value.uri + local_dir = ctx.file_access.get_random_local_directory() + ctx.file_access.get_data(path, local_dir, is_multipart=True) + return pd.read_parquet(local_dir) + + +class ArrowToParquetEncodingHandler(StructuredDatasetEncoder): + def __init__(self, protocol: str): + super().__init__(pa.Table, protocol, PARQUET) + + def encode( + self, + ctx: FlyteContext, + structured_dataset: StructuredDataset, + structured_dataset_type: StructuredDatasetType, + ) -> literals.StructuredDataset: + path = typing.cast(str, structured_dataset.uri) or ctx.file_access.get_random_remote_path() + df = structured_dataset.dataframe + local_dir = ctx.file_access.get_random_local_directory() + local_path = os.path.join(local_dir, f"{0:05}") + pq.write_table(df, local_path) + ctx.file_access.upload_directory(local_dir, path) + return literals.StructuredDataset(uri=path, metadata=StructuredDatasetMetadata(structured_dataset_type)) + + +class ParquetToArrowDecodingHandler(StructuredDatasetDecoder): + def __init__(self, protocol: str): + super().__init__(pa.Table, protocol, PARQUET) + + def decode( + self, + ctx: FlyteContext, + flyte_value: literals.StructuredDataset, + ) -> pa.Table: + path = flyte_value.uri + local_dir = ctx.file_access.get_random_local_directory() + ctx.file_access.get_data(path, local_dir, is_multipart=True) + return pq.read_table(local_dir) + + +for protocol in [LOCAL, S3]: # Should we add GCS + FLYTE_DATASET_TRANSFORMER.register_handler(PandasToParquetEncodingHandler(protocol), default_for_type=True) + FLYTE_DATASET_TRANSFORMER.register_handler(ParquetToPandasDecodingHandler(protocol), default_for_type=True) + FLYTE_DATASET_TRANSFORMER.register_handler(ArrowToParquetEncodingHandler(protocol), default_for_type=True) + FLYTE_DATASET_TRANSFORMER.register_handler(ParquetToArrowDecodingHandler(protocol), default_for_type=True) diff --git a/flytekit/types/structured/bigquery.py b/flytekit/types/structured/bigquery.py new file mode 100644 index 0000000000..33980a3e14 --- /dev/null +++ b/flytekit/types/structured/bigquery.py @@ -0,0 +1,112 @@ +import re +import typing + +import pandas as pd +import pyarrow as pa +from google.cloud import bigquery, bigquery_storage +from google.cloud.bigquery_storage_v1 import types + +from flytekit import FlyteContext +from flytekit.models import literals +from flytekit.models.types import StructuredDatasetType +from flytekit.types.structured.structured_dataset import ( + BIGQUERY, + DF, + FLYTE_DATASET_TRANSFORMER, + StructuredDataset, + StructuredDatasetDecoder, + StructuredDatasetEncoder, + StructuredDatasetMetadata, +) + + +def _write_to_bq(structured_dataset: StructuredDataset): + table_id = typing.cast(str, structured_dataset.uri).split("://", 1)[1].replace(":", ".") + client = bigquery.Client() + df = structured_dataset.dataframe + if isinstance(df, pa.Table): + df = df.to_pandas() + client.load_table_from_dataframe(df, table_id) + + +def _read_from_bq(flyte_value: literals.StructuredDataset) -> pd.DataFrame: + path = flyte_value.uri + _, project_id, dataset_id, table_id = re.split("\\.|://|:", path) + client = bigquery_storage.BigQueryReadClient() + table = f"projects/{project_id}/datasets/{dataset_id}/tables/{table_id}" + parent = "projects/{}".format(project_id) + + requested_session = types.ReadSession( + table=table, + data_format=types.DataFormat.ARROW, + ) + read_session = client.create_read_session(parent=parent, read_session=requested_session) + + stream = read_session.streams[0] + reader = client.read_rows(stream.name) + frames = [] + for message in reader.rows().pages: + frames.append(message.to_dataframe()) + return pd.concat(frames) + + +class PandasToBQEncodingHandlers(StructuredDatasetEncoder): + def __init__(self): + super().__init__(pd.DataFrame, BIGQUERY, supported_format="") + + def encode( + self, + ctx: FlyteContext, + structured_dataset: StructuredDataset, + structured_dataset_type: StructuredDatasetType, + ) -> literals.StructuredDataset: + _write_to_bq(structured_dataset) + return literals.StructuredDataset( + uri=typing.cast(str, structured_dataset.uri), metadata=StructuredDatasetMetadata(structured_dataset_type) + ) + + +class BQToPandasDecodingHandler(StructuredDatasetDecoder): + def __init__(self): + super().__init__(pd.DataFrame, BIGQUERY, supported_format="") + + def decode( + self, + ctx: FlyteContext, + flyte_value: literals.StructuredDataset, + ) -> typing.Union[DF, typing.Generator[DF, None, None]]: + return _read_from_bq(flyte_value) + + +class ArrowToBQEncodingHandlers(StructuredDatasetEncoder): + def __init__(self): + super().__init__(pa.Table, BIGQUERY, supported_format="") + + def encode( + self, + ctx: FlyteContext, + structured_dataset: StructuredDataset, + structured_dataset_type: StructuredDatasetType, + ) -> literals.StructuredDataset: + _write_to_bq(structured_dataset) + return literals.StructuredDataset( + uri=typing.cast(str, structured_dataset.uri), metadata=StructuredDatasetMetadata(structured_dataset_type) + ) + + +class BQToArrowDecodingHandler(StructuredDatasetDecoder): + def __init__(self): + super().__init__(pa.Table, BIGQUERY, supported_format="") + + def decode( + self, + ctx: FlyteContext, + flyte_value: literals.StructuredDataset, + ) -> typing.Union[DF, typing.Generator[DF, None, None]]: + return pa.Table.from_pandas(_read_from_bq(flyte_value)) + + +FLYTE_DATASET_TRANSFORMER.register_handler(PandasToBQEncodingHandlers(), default_for_type=False) +FLYTE_DATASET_TRANSFORMER.register_handler(BQToPandasDecodingHandler(), default_for_type=False) +FLYTE_DATASET_TRANSFORMER.register_handler(ArrowToBQEncodingHandlers(), default_for_type=False) +FLYTE_DATASET_TRANSFORMER.register_handler(BQToArrowDecodingHandler(), default_for_type=False) diff --git a/flytekit/types/structured/structured_dataset.py b/flytekit/types/structured/structured_dataset.py new file mode 100644 index 0000000000..fb9ab4f79f --- /dev/null +++ b/flytekit/types/structured/structured_dataset.py @@ -0,0 +1,584 @@ +from __future__ import annotations + +import collections +import os +import re +import types +import typing +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from typing import Dict, Generator, Optional, Type, Union + +from dataclasses_json import config, dataclass_json +from marshmallow import fields + +try: + from typing import Annotated, get_args, get_origin +except ImportError: + from typing_extensions import Annotated, get_origin, get_args + +import _datetime +import numpy as _np +import pyarrow as pa + +from flytekit.core.context_manager import FlyteContext, FlyteContextManager +from flytekit.core.type_engine import TypeTransformer +from flytekit.extend import TypeEngine +from flytekit.loggers import logger +from flytekit.models import literals +from flytekit.models import types as type_models +from flytekit.models.literals import Literal, Scalar, StructuredDatasetMetadata +from flytekit.models.types import LiteralType, StructuredDatasetType + +T = typing.TypeVar("T") # StructuredDataset type or a dataframe type +DF = typing.TypeVar("DF") # Dataframe type + +# Protocols +BIGQUERY = "bq" +S3 = "s3" +LOCAL = "/" + +# Storage formats +PARQUET = "parquet" + + +@dataclass_json +@dataclass +class StructuredDataset(object): + uri: typing.Optional[os.PathLike] = field(default=None, metadata=config(mm_field=fields.String())) + file_format: typing.Optional[str] = field(default=PARQUET, metadata=config(mm_field=fields.String())) + """ + This is the user facing StructuredDataset class. Please don't confuse it with the literals.StructuredDataset + class (that is just a model, a Python class representation of the protobuf). + """ + + FILE_FORMAT = PARQUET + + @classmethod + def columns(cls) -> typing.Dict[str, typing.Type]: + return {} + + @classmethod + def column_names(cls) -> typing.List[str]: + return [k for k, v in cls.columns().items()] + + def __class_getitem__(cls, args: typing.Union[typing.Dict[str, typing.Type], tuple]) -> Type[StructuredDataset]: + if args is None: + return cls + + format = PARQUET + if isinstance(args, tuple): + columns = args[0] + format = args[1] + else: + columns = args + + if not isinstance(columns, dict): + raise AssertionError( + f"Columns should be specified as an ordered dict " + f"of column names and their types, received {type(columns)}" + ) + + if not isinstance(format, str): + raise AssertionError(f"format should be specified as an string, received {type(format)}") + + # If nothing happened, columns and format are the default, then just use the main class + if len(columns) == 0 and format == PARQUET: + return cls + + class _TypedStructuredDataset(StructuredDataset): + # Get the type engine to see this as kind of a generic + __origin__ = StructuredDataset + FILE_FORMAT = format + + @classmethod + def columns(cls) -> typing.Dict[str, typing.Type]: + return columns + + return _TypedStructuredDataset + + def __init__( + self, + dataframe: typing.Optional[typing.Any] = None, + uri: Optional[str] = None, + metadata: typing.Optional[literals.StructuredDatasetMetadata] = None, + **kwargs, + ): + self._dataframe = dataframe + # Make these fields public, so that the dataclass transformer can set a value for it + # https://github.com/flyteorg/flytekit/blob/bcc8541bd6227b532f8462563fe8aac902242b21/flytekit/core/type_engine.py#L298 + self.uri = uri + # This is a special attribute that indicates if the data was either downloaded or uploaded + self._metadata = metadata + # This is not for users to set, the transformer will set this. + self._literal_sd: Optional[literals.StructuredDataset] = None + # Not meant for users to set, will be set by an open() call + self._dataframe_type = None + + @property + def dataframe(self) -> Type[typing.Any]: + return self._dataframe + + @property + def metadata(self) -> Optional[StructuredDatasetMetadata]: + return self._metadata + + @property + def literal(self) -> Optional[literals.StructuredDataset]: + return self._literal_sd + + def open(self, dataframe_type: Type[DF]): + self._dataframe_type = dataframe_type + return self + + def all(self) -> DF: + if self._dataframe_type is None: + raise ValueError("No dataframe type set. Use open() to set the local dataframe type you want to use.") + ctx = FlyteContextManager.current_context() + return FLYTE_DATASET_TRANSFORMER.open_as(ctx, self.literal, self._dataframe_type) + + def iter(self) -> Generator[DF, None, None]: + if self._dataframe_type is None: + raise ValueError("No dataframe type set. Use open() to set the local dataframe type you want to use.") + ctx = FlyteContextManager.current_context() + return FLYTE_DATASET_TRANSFORMER.iter_as(ctx, self.literal, self._dataframe_type) + + +class StructuredDatasetEncoder(ABC): + def __init__(self, python_type: Type[T], protocol: str, supported_format: Optional[str] = None): + """ + Extend this abstract class, implement the encode function, and register your concrete class with the + FLYTE_DATASET_TRANSFORMER defined at this module level in order for the core flytekit type engine to handle + dataframe libraries. This is the encoding interface, meaning it is used when there is a Python value that the + flytekit type engine is trying to convert into a Flyte Literal. For the other way, see + the StructuredDatasetEncoder + + :param python_type: The dataframe class in question that you want to register this encoder with + :param protocol: A prefix representing the storage driver (e.g. 's3, 'gs', 'bq', etc.). You can use either + "s3" or "s3://". They are the same since the "://" will just be stripped by the constructor. + :param supported_format: Arbitrary string representing the format. If not supplied then an empty string + will be used. An empty string implies that the encoder works with any format. If the format being asked + for does not exist, the transformer enginer will look for the "" endcoder instead and write a warning. + """ + self._python_type = python_type + self._protocol = protocol.replace("://", "") + self._supported_format = supported_format or "" + + @property + def python_type(self) -> Type[T]: + return self._python_type + + @property + def protocol(self) -> str: + return self._protocol + + @property + def supported_format(self) -> str: + return self._supported_format + + @abstractmethod + def encode( + self, + ctx: FlyteContext, + structured_dataset: StructuredDataset, + structured_dataset_type: StructuredDatasetType, + ) -> literals.StructuredDataset: + """ + Even if the user code returns a plain dataframe instance, the dataset transformer engine will wrap the + incoming dataframe with defaults set for that dataframe + type. This simplifies this function's interface as a lot of data that could be specified by the user using + the + # TODO: Do we need to add a flag to indicate if it was wrapped by the transformer or by the user? + + :param ctx: + :param structured_dataset: This is a StructuredDataset wrapper object. See more info above. + :param structured_dataset_type: This the StructuredDatasetType, as found in the LiteralType of the interface + of the task that invoked this encoding call. It is passed along to encoders so that authors of encoders + can include it in the returned literals.StructuredDataset. See the IDL for more information on why this + literal in particular carries the type information along with it. If the encoder doesn't supply it, it will + also be filled in after the encoder runs by the transformer engine. + :return: This function should return a StructuredDataset literal object. Do not confuse this with the + StructuredDataset wrapper class used as input to this function - that is the user facing Python class. + This function needs to return the IDL StructuredDataset. + """ + raise NotImplementedError + + +class StructuredDatasetDecoder(ABC): + def __init__(self, python_type: Type[DF], protocol: str, supported_format: Optional[str] = None): + """ + Extend this abstract class, implement the decode function, and register your concrete class with the + FLYTE_DATASET_TRANSFORMER defined at this module level in order for the core flytekit type engine to handle + dataframe libraries. This is the decoder interface, meaning it is used when there is a Flyte Literal value, + and we have to get a Python value out of it. For the other way, see the StructuredDatasetEncoder + + :param python_type: The dataframe class in question that you want to register this decoder with + :param protocol: A prefix representing the storage driver (e.g. 's3, 'gs', 'bq', etc.). You can use either + "s3" or "s3://". They are the same since the "://" will just be stripped by the constructor. + :param supported_format: Arbitrary string representing the format. If not supplied then an empty string + will be used. An empty string implies that the decoder works with any format. If the format being asked + for does not exist, the transformer enginer will look for the "" decoder instead and write a warning. + """ + self._python_type = python_type + self._protocol = protocol.replace("://", "") + self._supported_format = supported_format or "" + + @property + def python_type(self) -> Type[DF]: + return self._python_type + + @property + def protocol(self) -> str: + return self._protocol + + @property + def supported_format(self) -> str: + return self._supported_format + + @abstractmethod + def decode( + self, + ctx: FlyteContext, + flyte_value: literals.StructuredDataset, + ) -> Union[DF, Generator[DF, None, None]]: + """ + This is code that will be called by the dataset transformer engine to ultimately translate from a Flyte Literal + value into a Python instance. + + :param ctx: + :param flyte_value: This will be a Flyte IDL StructuredDataset Literal - do not confuse this with the + StructuredDataset class defined also in this module. + :return: This function can either return an instance of the dataframe that this decoder handles, or an iterator + of those dataframes. + """ + raise NotImplementedError + + +def protocol_prefix(uri: str) -> str: + g = re.search(r"([\w]+)://.*", uri) + if g and g.groups(): + return g.groups()[0] + return LOCAL + + +class StructuredDatasetTransformerEngine(TypeTransformer[StructuredDataset]): + """ + Think of this transformer as a higher-level meta transformer that is used for all the dataframe types. + If you are bringing a custom data frame type, or any data frame type, to flytekit, instead of + registering with the main type engine, you should register with this transformer instead. + """ + + _SUPPORTED_TYPES: typing.Dict[Type, LiteralType] = { + _np.int32: type_models.LiteralType(simple=type_models.SimpleType.INTEGER), + _np.int64: type_models.LiteralType(simple=type_models.SimpleType.INTEGER), + _np.uint32: type_models.LiteralType(simple=type_models.SimpleType.INTEGER), + _np.uint64: type_models.LiteralType(simple=type_models.SimpleType.INTEGER), + int: type_models.LiteralType(simple=type_models.SimpleType.INTEGER), + _np.float32: type_models.LiteralType(simple=type_models.SimpleType.FLOAT), + _np.float64: type_models.LiteralType(simple=type_models.SimpleType.FLOAT), + float: type_models.LiteralType(simple=type_models.SimpleType.FLOAT), + _np.bool_: type_models.LiteralType(simple=type_models.SimpleType.BOOLEAN), # type: ignore + bool: type_models.LiteralType(simple=type_models.SimpleType.BOOLEAN), + _np.datetime64: type_models.LiteralType(simple=type_models.SimpleType.DATETIME), + _datetime.datetime: type_models.LiteralType(simple=type_models.SimpleType.DATETIME), + _np.timedelta64: type_models.LiteralType(simple=type_models.SimpleType.DURATION), + _datetime.timedelta: type_models.LiteralType(simple=type_models.SimpleType.DURATION), + _np.string_: type_models.LiteralType(simple=type_models.SimpleType.STRING), + _np.str_: type_models.LiteralType(simple=type_models.SimpleType.STRING), + _np.object_: type_models.LiteralType(simple=type_models.SimpleType.STRING), + str: type_models.LiteralType(simple=type_models.SimpleType.STRING), + } + + ENCODERS: Dict[Type, Dict[str, Dict[str, StructuredDatasetEncoder]]] = {} + DECODERS: Dict[Type, Dict[str, Dict[str, StructuredDatasetDecoder]]] = {} + DEFAULT_PROTOCOLS: Dict[Type, str] = {} + DEFAULT_FORMATS: Dict[Type, str] = {} + + Handlers = Union[StructuredDatasetEncoder, StructuredDatasetDecoder] + + def _finder(self, handler_map, df_type: Type, protocol: str, format: str): + try: + return handler_map[df_type][protocol][format] + except KeyError: + try: + hh = handler_map[df_type][protocol][""] + logger.info( + f"Didn't find format specific handler {type(handler_map)} for protocol {protocol}" + f" format {format}, using default instead." + ) + return hh + except KeyError: + ... + raise ValueError(f"Failed to find a handler for {df_type}, protocol {protocol}, fmt {format}") + + def get_encoder(self, df_type: Type, protocol: str, format: str): + return self._finder(self.ENCODERS, df_type, protocol, format) + + def get_decoder(self, df_type: Type, protocol: str, format: str): + return self._finder(self.DECODERS, df_type, protocol, format) + + def _handler_finder(self, h: Handlers) -> Dict[str, Handlers]: + # Maybe think about default dict in the future, but is typing as nice? + if isinstance(h, StructuredDatasetEncoder): + top_level = self.ENCODERS + elif isinstance(h, StructuredDatasetDecoder): + top_level = self.DECODERS + else: + raise TypeError(f"We don't support this type of handler {h}") + if h.python_type not in top_level: + top_level[h.python_type] = {} + if h.protocol not in top_level[h.python_type]: + top_level[h.python_type][h.protocol] = {} + return top_level[h.python_type][h.protocol] + + def __init__(self): + super().__init__("StructuredDataset Transformer", StructuredDataset) + self._type_assertions_enabled = False + + def register_handler(self, h: Handlers, default_for_type: Optional[bool] = True, override: Optional[bool] = False): + """ + Call this with any handler to register it with this dataframe meta-transformer + + The string "://" should not be present in any handler's protocol so we don't check for it. + """ + lowest_level = self._handler_finder(h) + if h.supported_format in lowest_level and override is False: + raise ValueError(f"Already registered a handler for {(h.python_type, h.protocol, h.supported_format)}") + lowest_level[h.supported_format] = h + logger.debug(f"Registered {h} as handler for {h.python_type}, protocol {h.protocol}, fmt {h.supported_format}") + + if default_for_type: + # TODO: Add logging, think about better ux, maybe default False and warn if doesn't exist. + self.DEFAULT_FORMATS[h.python_type] = h.supported_format + self.DEFAULT_PROTOCOLS[h.python_type] = h.protocol + + # Register with the type engine as well + # The semantics as of now are such that it doesn't matter which order these transformers are loaded in, as + # long as the older Pandas/FlyteSchema transformer do not also specify the override + TypeEngine.register_additional_type(self, h.python_type, override=True) + + def assert_type(self, t: Type[StructuredDataset], v: typing.Any): + return + + def to_literal( + self, + ctx: FlyteContext, + python_val: Union[StructuredDataset, typing.Any], + python_type: Union[Type[StructuredDataset], Type], + expected: LiteralType, + ) -> Literal: + # Make a copy in case we need to hand off to encoders, since we can't be sure of mutations. + # Check first to see if it's even an SD type. For backwards compatibility, we may be getting a + if get_origin(python_type) is Annotated: + python_type = get_args(python_type)[0] + sdt = StructuredDatasetType(format=self.DEFAULT_FORMATS.get(python_type, None)) + + if expected and expected.structured_dataset_type: + sdt = StructuredDatasetType( + columns=expected.structured_dataset_type.columns, + format=expected.structured_dataset_type.format, + external_schema_type=expected.structured_dataset_type.external_schema_type, + external_schema_bytes=expected.structured_dataset_type.external_schema_bytes, + ) + + # If the type signature has the StructuredDataset class, it will, or at least should, also be a + # StructuredDataset instance. + if issubclass(python_type, StructuredDataset): + assert isinstance(python_val, StructuredDataset) + # There are three cases that we need to take care of here. + + # 1. A task returns a StructuredDataset that was just a passthrough input. If this happens + # then return the original literals.StructuredDataset without invoking any encoder + # + # Ex. + # def t1(dataset: StructuredDataset[my_cols]) -> StructuredDataset[my_cols]: + # return dataset + if python_val._literal_sd is not None: + if python_val.dataframe is not None: + raise ValueError( + f"Shouldn't have specified both literal {python_val._literal_sd} and dataframe {python_val.dataframe}" + ) + return Literal(scalar=Scalar(structured_dataset=python_val._literal_sd)) + + # 2. A task returns a python StructuredDataset with a uri. + # Note: this case is also what happens we start a local execution of a task with a python StructuredDataset. + # It gets converted into a literal first, then back into a python StructuredDataset. + # + # Ex. + # def t2(uri: str) -> StructuredDataset[my_cols] + # return StructuredDataset(uri=uri) + if python_val.dataframe is None: + if not python_val.uri: + raise ValueError(f"If dataframe is not specified, then the uri should be specified. {python_val}") + sd_model = literals.StructuredDataset( + uri=python_val.uri, + metadata=StructuredDatasetMetadata(structured_dataset_type=sdt), + ) + return Literal(scalar=Scalar(structured_dataset=sd_model)) + + # 3. This is the third and probably most common case. The python StructuredDataset object wraps a dataframe + # that we will need to invoke an encoder for. Figure out which encoder to call and invoke it. + df_type = type(python_val.dataframe) + if python_val.uri is None: + protocol = self.DEFAULT_PROTOCOLS[df_type] + else: + protocol = protocol_prefix(python_val.uri) + return self.encode( + ctx, + python_val, + df_type, + protocol, + python_val.file_format, + sdt, + ) + + # Otherwise assume it's a dataframe instance. Wrap it with some defaults + fmt = self.DEFAULT_FORMATS[python_type] + protocol = self.DEFAULT_PROTOCOLS[python_type] + meta = StructuredDatasetMetadata(structured_dataset_type=expected.structured_dataset_type if expected else None) + sd = StructuredDataset(dataframe=python_val, metadata=meta) + return self.encode(ctx, sd, python_type, protocol, fmt, sdt) + + def encode( + self, + ctx: FlyteContext, + sd: StructuredDataset, + df_type: Type, + protocol: str, + format: str, + structured_literal_type: StructuredDatasetType, + ) -> Literal: + handler: StructuredDatasetEncoder + handler = self.get_encoder(df_type, protocol, format) + sd_model = handler.encode(ctx, sd, structured_literal_type) + # This block is here in case the encoder did not set the type information in the metadata. Since this literal + # is special in that it carries around the type itself, we want to make sure the type info therein is at + # least as good as the type of the interface. + if sd_model.metadata is None: + sd_model._metadata = StructuredDatasetMetadata(structured_literal_type) + if sd_model.metadata.structured_dataset_type is None: + sd_model.metadata._structured_dataset_type = structured_literal_type + # Always set the format here to the format of the handler. + # Note that this will always be the same as the incoming format except for when the fallback handler + # with a format of "" is used. + sd_model.metadata._structured_dataset_type.format = handler.supported_format + return Literal(scalar=Scalar(structured_dataset=sd_model)) + + def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[T]) -> T: + # The literal that we get in might be an old FlyteSchema. + # We'll continue to support this for the time being. + if get_origin(expected_python_type) is Annotated: + expected_python_type = get_args(expected_python_type)[0] + if lv.scalar.schema is not None: + sd = StructuredDataset() + sd_literal = literals.StructuredDataset( + uri=lv.scalar.schema.uri, + metadata=literals.StructuredDatasetMetadata( + # Dataframe will always be serialized to parquet file by FlyteSchema transformer + structured_dataset_type=StructuredDatasetType(format=PARQUET) + ), + ) + sd._literal_sd = sd_literal + if issubclass(expected_python_type, StructuredDataset): + return sd + else: + return self.open_as(ctx, sd_literal, df_type=expected_python_type) + + # Either a StructuredDataset type or some dataframe type. + if issubclass(expected_python_type, StructuredDataset): + # Just save the literal for now. If in the future we find that we need the StructuredDataset type hint + # type also, we can add it. + sd = expected_python_type( + dataframe=None, + # Specifying these two are just done for completeness. Kind of waste since + # we're saving the whole incoming literal to _literal_sd. + metadata=lv.scalar.structured_dataset.metadata, + ) + sd._literal_sd = lv.scalar.structured_dataset + return sd + + # If the requested type was not a StructuredDataset, then it means it was a plain dataframe type, which means + # we should do the opening/downloading and whatever else it might entail right now. No iteration option here. + return self.open_as(ctx, lv.scalar.structured_dataset, df_type=expected_python_type) + + def open_as(self, ctx: FlyteContext, sd: literals.StructuredDataset, df_type: Type[DF]) -> DF: + protocol = protocol_prefix(sd.uri) + decoder = self.get_decoder(df_type, protocol, sd.metadata.structured_dataset_type.format) + result = decoder.decode(ctx, sd) + if isinstance(result, types.GeneratorType): + raise ValueError(f"Decoder {decoder} returned iterator {result} but whole value requested from {sd}") + return result + + def iter_as( + self, ctx: FlyteContext, sd: literals.StructuredDataset, df_type: Type[DF] + ) -> Generator[DF, None, None]: + protocol = protocol_prefix(sd.uri) + decoder = self.DECODERS[df_type][protocol][sd.metadata.structured_dataset_type.format] + result = decoder.decode(ctx, sd) + if not isinstance(result, types.GeneratorType): + raise ValueError(f"Decoder {decoder} didn't return iterator {result} but should have from {sd}") + return result + + def _get_dataset_column_literal_type(self, t: Type): + if t in self._SUPPORTED_TYPES: + return self._SUPPORTED_TYPES[t] + if hasattr(t, "__origin__") and t.__origin__ == list: + return type_models.LiteralType(collection_type=self._get_dataset_column_literal_type(t.__args__[0])) + if hasattr(t, "__origin__") and t.__origin__ == dict: + return type_models.LiteralType(map_value_type=self._get_dataset_column_literal_type(t.__args__[1])) + raise AssertionError(f"type {t} is currently not supported by StructuredDataset") + + def _get_dataset_type(self, t: typing.Union[Type[StructuredDataset], typing.Any]) -> StructuredDatasetType: + converted_cols: typing.List[StructuredDatasetType.DatasetColumn] = [] + # Handle different kinds of annotation + # my_cols = kwtypes(x=int, y=str) + # 1. Fill in format correctly by checking for typing.annotated. For example, Annotated[pd.Dataframe, my_cols] + if get_origin(t) is Annotated: + _, *hint_args = get_args(t) + if type(hint_args[0]) is collections.OrderedDict: + for k, v in hint_args[0].items(): + lt = self._get_dataset_column_literal_type(v) + converted_cols.append(StructuredDatasetType.DatasetColumn(name=k, literal_type=lt)) + return StructuredDatasetType(columns=converted_cols, format=PARQUET) + # 3. Fill in external schema type and bytes by checking for typing.annotated metadata. + # For example, Annotated[pd.Dataframe, pa.schema([("col1", pa.int32()), ("col2", pa.string())])] + elif type(hint_args[0]) is pa.lib.Schema: + return StructuredDatasetType( + format=PARQUET, + external_schema_type="arrow", + external_schema_bytes=typing.cast(pa.lib.Schema, hint_args[0]).to_string().encode(), + ) + raise ValueError(f"Unrecognized Annotated type for StructuredDataset {t}") + + # 2. Fill in columns by checking for StructuredDataset metadata. For example, StructuredDataset[my_cols, parquet] + elif issubclass(t, StructuredDataset): + for k, v in t.columns().items(): + lt = self._get_dataset_column_literal_type(v) + converted_cols.append(StructuredDatasetType.DatasetColumn(name=k, literal_type=lt)) + return StructuredDatasetType(columns=converted_cols, format=t.FILE_FORMAT) + + # 3. pd.Dataframe + else: + fmt = self.DEFAULT_FORMATS.get(t, PARQUET) + return StructuredDatasetType(columns=converted_cols, format=fmt) + + def get_literal_type(self, t: typing.Union[Type[StructuredDataset], typing.Any]) -> LiteralType: + """ + Provide a concrete implementation so that writers of custom dataframe handlers since there's nothing that + special about the literal type. Any dataframe type will always be associated with the structured dataset type. + The other aspects of it - columns, external schema type, etc. can be read from associated metadata. + + :param t: The python dataframe type, which is mostly ignored. + """ + return LiteralType(structured_dataset_type=self._get_dataset_type(t)) + + def guess_python_type(self, literal_type: LiteralType) -> Type[T]: + # todo: technically we should return the dataframe type specified in the constructor, but to do that, + # we'd have to store that, which we don't do today. See possibly #1363 + if literal_type.structured_dataset_type is not None: + return StructuredDataset + raise ValueError(f"StructuredDatasetTransformerEngine cannot reverse {literal_type}") + + +FLYTE_DATASET_TRANSFORMER = StructuredDatasetTransformerEngine() +TypeEngine.register(FLYTE_DATASET_TRANSFORMER) diff --git a/plugins/flytekit-greatexpectations/flytekitplugins/great_expectations/schema.py b/plugins/flytekit-greatexpectations/flytekitplugins/great_expectations/schema.py index 9e0cd88305..2dafc26c53 100644 --- a/plugins/flytekit-greatexpectations/flytekitplugins/great_expectations/schema.py +++ b/plugins/flytekit-greatexpectations/flytekitplugins/great_expectations/schema.py @@ -126,7 +126,7 @@ def to_literal( raise TypeError(f"{datatype} is not a supported type") def _flyte_schema( - self, is_runtime: bool, ctx: FlyteContext, ge_conf: GreatExpectationsFlyteConfig, lv: Literal + self, is_runtime: bool, ctx: FlyteContext, ge_conf: GreatExpectationsFlyteConfig, uri: str ) -> (FlyteSchema, str): temp_dataset = "" @@ -136,7 +136,7 @@ def _flyte_schema( raise ValueError("local_file_path is missing!") # copy parquet file to user-given directory - ctx.file_access.get_data(lv.scalar.schema.uri, ge_conf.local_file_path, is_multipart=True) + ctx.file_access.get_data(uri, ge_conf.local_file_path, is_multipart=True) temp_dataset = os.path.basename(ge_conf.local_file_path) @@ -146,7 +146,7 @@ def downloader(x, y): return ( FlyteSchema( local_path=ctx.file_access.get_random_local_directory(), - remote_path=lv.scalar.schema.uri, + remote_path=uri, downloader=downloader, supported_mode=SchemaOpenMode.READ, ) @@ -187,7 +187,12 @@ def to_python_value( if not ( lv and lv.scalar - and ((lv.scalar.primitive and lv.scalar.primitive.string_value) or lv.scalar.schema or lv.scalar.blob) + and ( + (lv.scalar.primitive and lv.scalar.primitive.string_value) + or lv.scalar.schema + or lv.scalar.blob + or lv.scalar.structured_dataset + ) ): raise AssertionError("Can only validate a literal string/FlyteFile/FlyteSchema value") @@ -226,7 +231,14 @@ def to_python_value( # FlyteSchema if lv.scalar.schema: - return_dataset, temp_dataset = self._flyte_schema(is_runtime=is_runtime, ctx=ctx, ge_conf=ge_conf, lv=lv) + return_dataset, temp_dataset = self._flyte_schema( + is_runtime=is_runtime, ctx=ctx, ge_conf=ge_conf, uri=lv.scalar.schema.uri + ) + + if lv.scalar.structured_dataset: + return_dataset, temp_dataset = self._flyte_schema( + is_runtime=is_runtime, ctx=ctx, ge_conf=ge_conf, uri=lv.scalar.structured_dataset.uri + ) # FlyteFile if lv.scalar.blob: @@ -260,7 +272,7 @@ def to_python_value( if is_runtime and lv.scalar.primitive: final_batch_request["runtime_parameters"]["query"] = dataset - elif is_runtime and lv.scalar.schema: + elif is_runtime and (lv.scalar.schema or lv.scalar.structured_dataset): final_batch_request["runtime_parameters"]["batch_data"] = return_dataset else: raise AssertionError("Can only use runtime_parameters for query(str)/schema data") diff --git a/plugins/flytekit-spark/flytekitplugins/spark/__init__.py b/plugins/flytekit-spark/flytekitplugins/spark/__init__.py index 0057d4f325..145497b030 100644 --- a/plugins/flytekit-spark/flytekitplugins/spark/__init__.py +++ b/plugins/flytekit-spark/flytekitplugins/spark/__init__.py @@ -1,2 +1,2 @@ -from .schema import SparkDataFrameSchemaReader, SparkDataFrameSchemaWriter, SparkDataFrameTransformer +from .schema import ParquetToSparkDecodingHandler, SparkToParquetEncodingHandler from .task import Spark, new_spark_session diff --git a/plugins/flytekit-spark/flytekitplugins/spark/schema.py b/plugins/flytekit-spark/flytekitplugins/spark/schema.py index 1cae101295..72dad4fc92 100644 --- a/plugins/flytekit-spark/flytekitplugins/spark/schema.py +++ b/plugins/flytekit-spark/flytekitplugins/spark/schema.py @@ -2,12 +2,21 @@ from typing import Type import pyspark +from pyspark.sql.dataframe import DataFrame from flytekit import FlyteContext from flytekit.extend import T, TypeEngine, TypeTransformer -from flytekit.models.literals import Literal, Scalar, Schema -from flytekit.models.types import LiteralType, SchemaType +from flytekit.models import literals +from flytekit.models.literals import Literal, Scalar, Schema, StructuredDatasetMetadata +from flytekit.models.types import LiteralType, SchemaType, StructuredDatasetType from flytekit.types.schema import SchemaEngine, SchemaFormat, SchemaHandler, SchemaReader, SchemaWriter +from flytekit.types.structured.structured_dataset import ( + FLYTE_DATASET_TRANSFORMER, + PARQUET, + StructuredDataset, + StructuredDatasetDecoder, + StructuredDatasetEncoder, +) class SparkDataFrameSchemaReader(SchemaReader[pyspark.sql.DataFrame]): @@ -97,3 +106,37 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: # %% # This makes pyspark.DataFrame as a supported output/input type with flytekit. TypeEngine.register(SparkDataFrameTransformer()) + + +class SparkToParquetEncodingHandler(StructuredDatasetEncoder): + def __init__(self, protocol: str): + super().__init__(DataFrame, protocol, PARQUET) + + def encode( + self, + ctx: FlyteContext, + structured_dataset: StructuredDataset, + structured_dataset_type: StructuredDatasetType, + ) -> literals.StructuredDataset: + path = typing.cast(str, structured_dataset.uri) or ctx.file_access.get_random_remote_directory() + df = typing.cast(DataFrame, structured_dataset.dataframe) + df.write.mode("overwrite").parquet(path) + return literals.StructuredDataset(uri=path, metadata=StructuredDatasetMetadata(structured_dataset_type)) + + +class ParquetToSparkDecodingHandler(StructuredDatasetDecoder): + def __init__(self, protocol: str): + super().__init__(DataFrame, protocol, PARQUET) + + def decode( + self, + ctx: FlyteContext, + flyte_value: literals.StructuredDataset, + ) -> DataFrame: + user_ctx = FlyteContext.current_context().user_space_params + return user_ctx.spark_session.read.parquet(flyte_value.uri) + + +for protocol in ["/", "s3"]: + FLYTE_DATASET_TRANSFORMER.register_handler(SparkToParquetEncodingHandler(protocol), default_for_type=True) + FLYTE_DATASET_TRANSFORMER.register_handler(ParquetToSparkDecodingHandler(protocol), default_for_type=True) diff --git a/plugins/flytekit-spark/tests/test_wf.py b/plugins/flytekit-spark/tests/test_wf.py index 56b7734f5f..a0a624fec7 100644 --- a/plugins/flytekit-spark/tests/test_wf.py +++ b/plugins/flytekit-spark/tests/test_wf.py @@ -1,4 +1,4 @@ -import pandas +import pandas as pd import pyspark from flytekitplugins.spark.task import Spark @@ -33,7 +33,7 @@ def test_spark_dataframe_input(): @task def my_dataset() -> my_schema: - return pandas.DataFrame(data={"name": ["Alice"], "age": [5]}) + return pd.DataFrame(data={"name": ["Alice"], "age": [5]}) @task(task_config=Spark()) def my_spark(df: pyspark.sql.DataFrame) -> my_schema: @@ -53,6 +53,54 @@ def my_wf() -> my_schema: assert df2 is not None +def test_ddwf1_with_spark(): + @task(task_config=Spark()) + def my_spark(a: int) -> (int, str): + session = flytekit.current_context().spark_session + assert session.sparkContext.appName == "FlyteSpark: ex:local:local:local" + return a + 2, "world" + + @task + def t2(a: str, b: str) -> str: + return b + a + + @workflow + def my_wf(a: int, b: str) -> (int, str): + x, y = my_spark(a=a) + d = t2(a=y, b=b) + return x, d + + x = my_wf(a=5, b="hello ") + assert x == (7, "hello world") + + +def test_fs_sd_compatibility(): + my_schema = FlyteSchema[kwtypes(name=str, age=int)] + + @task + def my_dataset() -> pd.DataFrame: + return pd.DataFrame(data={"name": ["Alice"], "age": [5]}) + + @task(task_config=Spark()) + def my_spark(df: pyspark.sql.DataFrame) -> my_schema: + session = flytekit.current_context().spark_session + new_df = session.createDataFrame([("Bob", 10)], my_schema.column_names()) + return df.union(new_df) + + @task(task_config=Spark()) + def read_spark_df(df: pyspark.sql.DataFrame) -> int: + return df.count() + + @workflow + def my_wf() -> int: + df = my_dataset() + fs = my_spark(df=df) + return read_spark_df(df=fs) + + res = my_wf() + assert res == 2 + + def test_spark_dataframe_return(): my_schema = FlyteSchema[kwtypes(name=str, age=int)] @@ -68,9 +116,7 @@ def my_wf(a: int) -> my_schema: return my_spark(a=a) x = my_wf(a=5) - reader = x.open(pandas.DataFrame) + reader = x.open(pd.DataFrame) df2 = reader.all() - result_df = df2.reset_index(drop=True) == pandas.DataFrame(data={"name": ["Alice"], "age": [5]}).reset_index( - drop=True - ) + result_df = df2.reset_index(drop=True) == pd.DataFrame(data={"name": ["Alice"], "age": [5]}).reset_index(drop=True) assert result_df.all().all() diff --git a/requirements-spark2.txt b/requirements-spark2.txt index a48ec4b5e2..ecc11d57ea 100644 --- a/requirements-spark2.txt +++ b/requirements-spark2.txt @@ -83,7 +83,7 @@ entrypoints==0.3 # jupyter-client # nbconvert # papermill -flyteidl==0.21.13 +flyteidl==0.21.17 # via flytekit gevent==21.12.0 # via sagemaker-training diff --git a/requirements.txt b/requirements.txt index 04cdd38a16..9ab1213845 100644 --- a/requirements.txt +++ b/requirements.txt @@ -81,7 +81,7 @@ entrypoints==0.3 # jupyter-client # nbconvert # papermill -flyteidl==0.21.13 +flyteidl==0.21.17 # via flytekit gevent==21.12.0 # via sagemaker-training diff --git a/setup.py b/setup.py index 5bf9f83aa7..bb5f9d1f1c 100644 --- a/setup.py +++ b/setup.py @@ -64,7 +64,7 @@ ] }, install_requires=[ - "flyteidl>=0.21.4", + "flyteidl>=0.21.17", "wheel>=0.30.0,<1.0.0", "pandas>=1.0.0,<2.0.0", "pyarrow>=4.0.0,<7.0.0", diff --git a/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt b/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt index 627e87980e..8fcc781de3 100644 --- a/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt +++ b/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt @@ -44,7 +44,7 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.21.13 +flyteidl==0.21.17 # via flytekit flytekit==0.25.0 # via -r tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.in diff --git a/tests/flytekit/unit/core/test_flyte_file.py b/tests/flytekit/unit/core/test_flyte_file.py index 05d8156ece..9132be082c 100644 --- a/tests/flytekit/unit/core/test_flyte_file.py +++ b/tests/flytekit/unit/core/test_flyte_file.py @@ -85,7 +85,7 @@ def my_wf() -> FlyteFile: ctx = context_manager.FlyteContext.current_context() with context_manager.FlyteContextManager.with_context(ctx.with_file_access(fs)): top_level_files = os.listdir(random_dir) - assert len(top_level_files) == 1 # the local_flytekit folder + assert len(top_level_files) == 1 # the flytekit_local folder x = my_wf() @@ -188,7 +188,6 @@ def my_wf() -> FlyteFile: # This creates a random directory that we know is empty. random_dir = context_manager.FlyteContext.current_context().file_access.get_random_local_directory() - print(f"Random {random_dir}") # Creating a new FileAccessProvider will add two folderst to the random dir fs = FileAccessProvider(local_sandbox_dir=random_dir, raw_output_prefix=os.path.join(random_dir, "mock_remote")) ctx = context_manager.FlyteContext.current_context() diff --git a/tests/flytekit/unit/core/test_imperative.py b/tests/flytekit/unit/core/test_imperative.py index fd45c831e8..6b99c93368 100644 --- a/tests/flytekit/unit/core/test_imperative.py +++ b/tests/flytekit/unit/core/test_imperative.py @@ -16,6 +16,7 @@ from flytekit.models import literals as literal_models from flytekit.types.file import FlyteFile from flytekit.types.schema import FlyteSchema +from flytekit.types.structured.structured_dataset import StructuredDatasetType default_img = Image(name="default", fqn="test", tag="tag") serialization_settings = context_manager.SerializationSettings( @@ -348,4 +349,7 @@ def ref_t2( assert wf_spec.template.interface.inputs["sqlite_archive"].type.blob is not None assert len(wf_spec.template.interface.outputs) == 1 - assert wf_spec.template.interface.outputs["output_from_t3"].type.schema is not None + assert wf_spec.template.interface.outputs["output_from_t3"].type.structured_dataset_type is not None + assert wf_spec.template.interface.outputs["output_from_t3"].type.structured_dataset_type == StructuredDatasetType( + format="parquet" + ) diff --git a/tests/flytekit/unit/core/test_interface.py b/tests/flytekit/unit/core/test_interface.py index 8e55ee1bb4..2317bab1c2 100644 --- a/tests/flytekit/unit/core/test_interface.py +++ b/tests/flytekit/unit/core/test_interface.py @@ -15,6 +15,11 @@ from flytekit.types.file import FlyteFile from flytekit.types.pickle import FlytePickle +try: + from typing import Annotated +except ImportError: + from typing_extensions import Annotated + def test_extract_only(): def x() -> typing.NamedTuple("NT1", x_str=str, y_int=int): @@ -184,6 +189,16 @@ def z(a: int = 7, b: str = "eleven") -> typing.Tuple[int, str]: assert not params.parameters["b"].required assert params.parameters["b"].default.scalar.primitive.string_value == "eleven" + def z(a: Annotated[int, "some annotation"]) -> Annotated[int, "some annotation"]: + return a + + our_interface = transform_function_to_interface(z) + params = transform_inputs_to_parameters(ctx, our_interface) + assert params.parameters["a"].required + assert params.parameters["a"].default is None + assert our_interface.inputs == {"a": Annotated[int, "some annotation"]} + assert our_interface.outputs == {"o0": Annotated[int, "some annotation"]} + def test_parameters_with_docstring(): ctx = context_manager.FlyteContext.current_context() diff --git a/tests/flytekit/unit/core/test_local_cache.py b/tests/flytekit/unit/core/test_local_cache.py index c43421bbf9..bf4e4c2edb 100644 --- a/tests/flytekit/unit/core/test_local_cache.py +++ b/tests/flytekit/unit/core/test_local_cache.py @@ -250,3 +250,15 @@ def my_wf(a: int, b: str) -> (int, str): x = my_wf(a=5, b="hello") assert x == (7, "hello world") assert n_cached_task_calls == 2 + + +""" +Update SD transformer so that it can to_python_value a Schema literal + - If a Schema literal is detected, copy the uri and use the new decoder to unwrap the uri +Update FS transformer so that it can to_python_value a StructuredDataset literal + - If a StructuredDataset literal is detected, use the uri from that instead. + +Update all plugins that can take in a FlyteSchema to also be able to take in a StructuredDataset. + +All tests should work with the presence of SD imports. +""" diff --git a/tests/flytekit/unit/core/test_structured_dataset.py b/tests/flytekit/unit/core/test_structured_dataset.py new file mode 100644 index 0000000000..cb8ad85071 --- /dev/null +++ b/tests/flytekit/unit/core/test_structured_dataset.py @@ -0,0 +1,253 @@ +import typing + +import pytest + +from flytekit.core import context_manager +from flytekit.core.context_manager import FlyteContext, FlyteContextManager, Image, ImageConfig +from flytekit.core.type_engine import TypeEngine +from flytekit.models import literals +from flytekit.models.literals import StructuredDatasetMetadata +from flytekit.models.types import SimpleType, StructuredDatasetType + +try: + from typing import Annotated +except ImportError: + from typing_extensions import Annotated + +import pandas as pd +import pyarrow as pa + +from flytekit import kwtypes +from flytekit.types.structured.structured_dataset import ( + FLYTE_DATASET_TRANSFORMER, + PARQUET, + StructuredDataset, + StructuredDatasetDecoder, + StructuredDatasetEncoder, + protocol_prefix, +) + +my_cols = kwtypes(w=typing.Dict[str, typing.Dict[str, int]], x=typing.List[typing.List[int]], y=int, z=str) + +fields = [("some_int", pa.int32()), ("some_string", pa.string())] +arrow_schema = pa.schema(fields) + +serialization_settings = context_manager.SerializationSettings( + project="proj", + domain="dom", + version="123", + image_config=ImageConfig(Image(name="name", fqn="asdf/fdsa", tag="123")), + env={}, +) + + +def test_protocol(): + assert protocol_prefix("s3://my-s3-bucket/file") == "s3" + assert protocol_prefix("/file") == "/" + + +def generate_pandas() -> pd.DataFrame: + return pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]}) + + +def test_types_pandas(): + pt = pd.DataFrame + lt = TypeEngine.to_literal_type(pt) + assert lt.structured_dataset_type is not None + assert lt.structured_dataset_type.format == PARQUET + assert lt.structured_dataset_type.columns == [] + + +def test_types_annotated(): + pt = Annotated[pd.DataFrame, my_cols] + lt = TypeEngine.to_literal_type(pt) + assert len(lt.structured_dataset_type.columns) == 4 + assert lt.structured_dataset_type.columns[0].literal_type.map_value_type.map_value_type.simple == SimpleType.INTEGER + assert ( + lt.structured_dataset_type.columns[1].literal_type.collection_type.collection_type.simple == SimpleType.INTEGER + ) + assert lt.structured_dataset_type.columns[2].literal_type.simple == SimpleType.INTEGER + assert lt.structured_dataset_type.columns[3].literal_type.simple == SimpleType.STRING + + pt = Annotated[pd.DataFrame, arrow_schema] + lt = TypeEngine.to_literal_type(pt) + assert lt.structured_dataset_type.external_schema_type == "arrow" + assert "some_string" in str(lt.structured_dataset_type.external_schema_bytes) + + pt = Annotated[pd.DataFrame, kwtypes(a=None)] + with pytest.raises(AssertionError, match="type None is currently not supported by StructuredDataset"): + TypeEngine.to_literal_type(pt) + + pt = Annotated[pd.DataFrame, None] + with pytest.raises(ValueError, match="Unrecognized Annotated type for StructuredDataset"): + TypeEngine.to_literal_type(pt) + + +def test_types_sd(): + pt = StructuredDataset + lt = TypeEngine.to_literal_type(pt) + assert lt.structured_dataset_type is not None + + pt = StructuredDataset[my_cols] + lt = TypeEngine.to_literal_type(pt) + assert len(lt.structured_dataset_type.columns) == 4 + + pt = StructuredDataset[my_cols, "csv"] + lt = TypeEngine.to_literal_type(pt) + assert len(lt.structured_dataset_type.columns) == 4 + assert lt.structured_dataset_type.format == "csv" + + pt = StructuredDataset[{}, "csv"] + assert pt.FILE_FORMAT == "csv" + lt = TypeEngine.to_literal_type(pt) + assert len(lt.structured_dataset_type.columns) == 0 + assert lt.structured_dataset_type.format == "csv" + + +def test_retrieving(): + assert FLYTE_DATASET_TRANSFORMER.get_encoder(pd.DataFrame, "/", PARQUET) is not None + with pytest.raises(ValueError): + # We don't have a default "" format encoder + FLYTE_DATASET_TRANSFORMER.get_encoder(pd.DataFrame, "/", "") + + class TempEncoder(StructuredDatasetEncoder): + def __init__(self, protocol): + super().__init__(MyDF, protocol) + + def encode(self): + ... + + FLYTE_DATASET_TRANSFORMER.register_handler(TempEncoder("gs"), default_for_type=False) + with pytest.raises(ValueError): + FLYTE_DATASET_TRANSFORMER.register_handler(TempEncoder("gs://"), default_for_type=False) + + class TempEncoder: + pass + + with pytest.raises(TypeError, match="We don't support this type of handler"): + FLYTE_DATASET_TRANSFORMER.register_handler(TempEncoder, default_for_type=False) + + +def test_to_literal(): + ctx = FlyteContextManager.current_context() + lt = TypeEngine.to_literal_type(pd.DataFrame) + df = generate_pandas() + + lit = FLYTE_DATASET_TRANSFORMER.to_literal(ctx, df, python_type=pd.DataFrame, expected=lt) + assert lit.scalar.structured_dataset.metadata.structured_dataset_type.format == PARQUET + assert lit.scalar.structured_dataset.metadata.structured_dataset_type.format == PARQUET + + sd_with_literal_and_df = StructuredDataset(df) + sd_with_literal_and_df._literal_sd = lit + + with pytest.raises(ValueError, match="Shouldn't have specified both literal"): + FLYTE_DATASET_TRANSFORMER.to_literal(ctx, sd_with_literal_and_df, python_type=StructuredDataset, expected=lt) + + sd_with_nothing = StructuredDataset() + with pytest.raises(ValueError, match="If dataframe is not specified"): + FLYTE_DATASET_TRANSFORMER.to_literal(ctx, sd_with_nothing, python_type=StructuredDataset, expected=lt) + + sd_with_uri = StructuredDataset(uri="s3://some/extant/df.parquet") + + lt = TypeEngine.to_literal_type(StructuredDataset[{}, "new-df-format"]) + lit = FLYTE_DATASET_TRANSFORMER.to_literal(ctx, sd_with_uri, python_type=StructuredDataset, expected=lt) + assert lit.scalar.structured_dataset.uri == "s3://some/extant/df.parquet" + assert lit.scalar.structured_dataset.metadata.structured_dataset_type.format == "new-df-format" + + +class MyDF(pd.DataFrame): + ... + + +def test_fill_in_literal_type(): + class TempEncoder(StructuredDatasetEncoder): + def __init__(self, fmt: str): + super().__init__(MyDF, "tmpfs://", supported_format=fmt) + + def encode( + self, + ctx: FlyteContext, + structured_dataset: StructuredDataset, + structured_dataset_type: StructuredDatasetType, + ) -> literals.StructuredDataset: + return literals.StructuredDataset(uri="") + + FLYTE_DATASET_TRANSFORMER.register_handler(TempEncoder("myavro"), default_for_type=True) + lt = TypeEngine.to_literal_type(MyDF) + assert lt.structured_dataset_type.format == "myavro" + + ctx = FlyteContextManager.current_context() + sd = StructuredDataset(dataframe=42) + l = FLYTE_DATASET_TRANSFORMER.to_literal(ctx, sd, MyDF, lt) + # Test that the literal type is filled in even though the encode function above doesn't do it. + assert l.scalar.structured_dataset.metadata.structured_dataset_type.format == "myavro" + + # Test that looking up encoders/decoders falls back to the "" encoder/decoder + empty_format_temp_encoder = TempEncoder("") + FLYTE_DATASET_TRANSFORMER.register_handler(empty_format_temp_encoder, default_for_type=False) + + res = FLYTE_DATASET_TRANSFORMER.get_encoder(MyDF, "tmpfs", "rando") + assert res is empty_format_temp_encoder + + +def test_sd(): + sd = StructuredDataset(dataframe="hi") + sd.uri = "my uri" + assert sd.file_format == PARQUET + + with pytest.raises(ValueError, match="No dataframe type set"): + sd.all() + + with pytest.raises(ValueError, match="No dataframe type set."): + sd.iter() + + class MockPandasDecodingHandlers(StructuredDatasetDecoder): + def decode( + self, + ctx: FlyteContext, + flyte_value: literals.StructuredDataset, + ) -> typing.Union[typing.Generator[pd.DataFrame, None, None]]: + yield pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]}) + + FLYTE_DATASET_TRANSFORMER.register_handler( + MockPandasDecodingHandlers(pd.DataFrame, "tmpfs"), default_for_type=False + ) + sd = StructuredDataset() + sd._literal_sd = literals.StructuredDataset( + uri="tmpfs://somewhere", metadata=StructuredDatasetMetadata(StructuredDatasetType(format="")) + ) + assert isinstance(sd.open(pd.DataFrame).iter(), typing.Generator) + + with pytest.raises(ValueError): + sd.open(pd.DataFrame).all() + + class MockPandasDecodingHandlers(StructuredDatasetDecoder): + def decode( + self, + ctx: FlyteContext, + flyte_value: literals.StructuredDataset, + ) -> pd.DataFrame: + pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]}) + + FLYTE_DATASET_TRANSFORMER.register_handler( + MockPandasDecodingHandlers(pd.DataFrame, "tmpfs"), default_for_type=False, override=True + ) + sd = StructuredDataset() + sd._literal_sd = literals.StructuredDataset( + uri="tmpfs://somewhere", metadata=StructuredDatasetMetadata(StructuredDatasetType(format="")) + ) + + with pytest.raises(ValueError): + sd.open(pd.DataFrame).iter() + + +def test_class_getitem(): + assert StructuredDataset[None] == StructuredDataset + assert StructuredDataset[{}] == StructuredDataset + assert StructuredDataset[{"a": int}].FILE_FORMAT == StructuredDataset[{"a": int}, PARQUET].FILE_FORMAT + + with pytest.raises(AssertionError, match="Columns should be specified as an ordered dict"): + StructuredDataset[int] + + with pytest.raises(AssertionError, match="format should be specified as an string"): + StructuredDataset[{"a": int}, 123] diff --git a/tests/flytekit/unit/core/test_structured_dataset_handlers.py b/tests/flytekit/unit/core/test_structured_dataset_handlers.py new file mode 100644 index 0000000000..0d755ab78b --- /dev/null +++ b/tests/flytekit/unit/core/test_structured_dataset_handlers.py @@ -0,0 +1,43 @@ +import typing + +import pandas as pd +import pyarrow as pa +import pytest + +from flytekit.core import context_manager +from flytekit.core.base_task import kwtypes +from flytekit.models.types import StructuredDatasetType +from flytekit.types.structured import basic_dfs +from flytekit.types.structured.structured_dataset import ( + StructuredDataset, + StructuredDatasetDecoder, + StructuredDatasetEncoder, +) + +my_cols = kwtypes(w=typing.Dict[str, typing.Dict[str, int]], x=typing.List[typing.List[int]], y=int, z=str) + +fields = [("some_int", pa.int32()), ("some_string", pa.string())] +arrow_schema = pa.schema(fields) + + +def test_pandas(): + df = pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]}) + encoder = basic_dfs.PandasToParquetEncodingHandler("/") + decoder = basic_dfs.ParquetToPandasDecodingHandler("/") + + ctx = context_manager.FlyteContextManager.current_context() + sd = StructuredDataset( + dataframe=df, + ) + sd_lit = encoder.encode(ctx, sd, StructuredDatasetType(format="parquet")) + + df2 = decoder.decode(ctx, sd_lit) + assert df.equals(df2) + + +def test_base_isnt_instantiable(): + with pytest.raises(TypeError): + StructuredDatasetEncoder(pd.DataFrame, "", "") + + with pytest.raises(TypeError): + StructuredDatasetDecoder(pd.DataFrame, "", "") diff --git a/tests/flytekit/unit/core/test_type_delayed.py b/tests/flytekit/unit/core/test_type_delayed.py index 87daa91f47..268bf64285 100644 --- a/tests/flytekit/unit/core/test_type_delayed.py +++ b/tests/flytekit/unit/core/test_type_delayed.py @@ -5,8 +5,15 @@ from dataclasses_json import dataclass_json +from flytekit.core import context_manager +from flytekit.core.interface import transform_function_to_interface, transform_inputs_to_parameters from flytekit.core.type_engine import TypeEngine +try: + from typing import Annotated +except ImportError: + from typing_extensions import Annotated + @dataclass_json @dataclass @@ -25,3 +32,17 @@ def test_jsondc_schemaize(): # This test basically tests the broken behavior. Remove this test if # https://github.com/lovasoa/marshmallow_dataclass/issues/13 is ever fixed. assert pt is dict + + +def test_structured_dataset(): + ctx = context_manager.FlyteContext.current_context() + + def z(a: Annotated[int, "some annotation"]) -> Annotated[int, "some annotation"]: + return a + + our_interface = transform_function_to_interface(z) + params = transform_inputs_to_parameters(ctx, our_interface) + assert params.parameters["a"].required + assert params.parameters["a"].default is None + assert our_interface.inputs == {"a": Annotated[int, "some annotation"]} + assert our_interface.outputs == {"o0": Annotated[int, "some annotation"]} diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 0768b07db6..50ce9d99ac 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -7,6 +7,7 @@ from enum import Enum import pandas as pd +import pyarrow as pa import pytest from dataclasses_json import DataClassJsonMixin, dataclass_json from flyteidl.core import errors_pb2 @@ -14,6 +15,7 @@ from google.protobuf import struct_pb2 as _struct from marshmallow_enum import LoadDumpOptions from marshmallow_jsonschema import JSONSchema +from pandas._testing import assert_frame_equal import flytekit.common.exceptions.user as user_exceptions from flytekit import kwtypes @@ -43,6 +45,7 @@ from flytekit.types.pickle import FlytePickle from flytekit.types.pickle.pickle import FlytePickleTransformer from flytekit.types.schema import FlyteSchema +from flytekit.types.structured.structured_dataset import StructuredDataset T = typing.TypeVar("T") @@ -590,6 +593,35 @@ class TestFileStruct(object): assert o.b.c["hello"].path == ot.b.c["hello"].path +def test_structured_dataset_in_dataclass(): + df = pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]}) + + @dataclass_json + @dataclass + class InnerDatasetStruct(object): + a: StructuredDataset + + @dataclass_json + @dataclass + class DatasetStruct(object): + a: StructuredDataset + b: InnerDatasetStruct + + sd = StructuredDataset(dataframe=df, file_format="parquet") + o = DatasetStruct(a=sd, b=InnerDatasetStruct(a=sd)) + + ctx = FlyteContext.current_context() + tf = DataclassTransformer() + lt = tf.get_literal_type(DatasetStruct) + lv = tf.to_literal(ctx, o, DatasetStruct, lt) + ot = tf.to_python_value(ctx, lv=lv, expected_python_type=DatasetStruct) + + assert_frame_equal(df, ot.a.open(pd.DataFrame).all()) + assert_frame_equal(df, ot.b.a.open(pd.DataFrame).all()) + assert "parquet" == ot.a.file_format + assert "parquet" == ot.b.a.file_format + + # Enums should have string values class Color(Enum): RED = "red" @@ -604,6 +636,29 @@ class UnsupportedEnumValues(Enum): BLUE = 3 +def test_structured_dataset_type(): + name = "Name" + age = "Age" + data = {name: ["Tom", "Joseph"], age: [20, 22]} + df = pd.DataFrame(data) + + from flytekit.types.structured.structured_dataset import StructuredDataset, StructuredDatasetTransformerEngine + + tf = StructuredDatasetTransformerEngine() + lt = tf.get_literal_type(StructuredDataset[{name: str, age: int}, "parquet"]) + assert lt.structured_dataset_type is not None + + ctx = FlyteContextManager.current_context() + lv = tf.to_literal(ctx, df, pd.DataFrame, lt) + assert "/tmp/flyte" in lv.scalar.structured_dataset.uri + metadata = lv.scalar.structured_dataset.metadata + assert metadata.structured_dataset_type.format == "parquet" + v1 = tf.to_python_value(ctx, lv, pd.DataFrame) + v2 = tf.to_python_value(ctx, lv, pa.Table) + assert_frame_equal(df, v1) + assert_frame_equal(df, v2.to_pandas()) + + def test_enum_type(): t = TypeEngine.to_literal_type(Color) assert t is not None diff --git a/tests/flytekit/unit/core/test_type_hints.py b/tests/flytekit/unit/core/test_type_hints.py index 2f4a4e3068..bdb9e38c3d 100644 --- a/tests/flytekit/unit/core/test_type_hints.py +++ b/tests/flytekit/unit/core/test_type_hints.py @@ -11,9 +11,11 @@ from textwrap import dedent import pandas +import pandas as pd import pytest from dataclasses_json import dataclass_json from google.protobuf.struct_pb2 import Struct +from pandas._testing import assert_frame_equal import flytekit from flytekit import ContainerTask, Secret, SQLTask, dynamic, kwtypes, map_task @@ -38,6 +40,7 @@ from flytekit.types.directory import FlyteDirectory, TensorboardLogs from flytekit.types.file import FlyteFile, PNGImageFile from flytekit.types.schema import FlyteSchema, SchemaOpenMode +from flytekit.types.structured.structured_dataset import StructuredDataset serialization_settings = context_manager.SerializationSettings( project="proj", @@ -434,6 +437,36 @@ def wf(path: str) -> os.PathLike: assert "/tmp/flyte/" in wf(path="s3://somewhere").path +def test_structured_dataset_in_dataclass(): + df = pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]}) + + @dataclass_json + @dataclass + class InnerDatasetStruct(object): + a: StructuredDataset + + @dataclass_json + @dataclass + class DatasetStruct(object): + a: StructuredDataset + b: InnerDatasetStruct + + @task + def t1(path: str) -> DatasetStruct: + sd = StructuredDataset(dataframe=df, uri=path) + return DatasetStruct(a=sd, b=InnerDatasetStruct(a=sd)) + + @workflow + def wf(path: str) -> DatasetStruct: + return t1(path=path) + + res = wf(path="/tmp/somewhere") + assert "parquet" == res.a.file_format + assert "parquet" == res.b.a.file_format + assert_frame_equal(df, res.a.open(pd.DataFrame).all()) + assert_frame_equal(df, res.b.a.open(pd.DataFrame).all()) + + def test_wf1_with_map(): @task def t1(a: int) -> int: diff --git a/tests/flytekit/unit/core/test_workflows.py b/tests/flytekit/unit/core/test_workflows.py index df1d4a1e34..5054adc3b2 100644 --- a/tests/flytekit/unit/core/test_workflows.py +++ b/tests/flytekit/unit/core/test_workflows.py @@ -1,8 +1,11 @@ import typing from collections import OrderedDict +import pandas as pd import pytest +from pandas.testing import assert_frame_equal +from flytekit import StructuredDataset, kwtypes from flytekit.common.exceptions.user import FlyteValidationException, FlyteValueException from flytekit.common.translator import get_serializable from flytekit.core import context_manager @@ -10,6 +13,12 @@ from flytekit.core.context_manager import Image, ImageConfig from flytekit.core.task import task from flytekit.core.workflow import WorkflowFailurePolicy, WorkflowMetadata, WorkflowMetadataDefaults, workflow +from flytekit.types.schema import FlyteSchema + +try: + from typing import Annotated +except ImportError: + from typing_extensions import Annotated default_img = Image(name="default", fqn="test", tag="tag") serialization_settings = context_manager.SerializationSettings( @@ -256,3 +265,57 @@ def test_wf_docstring(): assert model_wf.template.interface.outputs["o1"].description == "outputs" assert len(model_wf.template.interface.inputs) == 1 assert model_wf.template.interface.inputs["a"].description == "input a" + + +my_cols = kwtypes(y=int, z=int) +pd_df = pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]}) + + +@task +def t1() -> Annotated[pd.DataFrame, my_cols]: + return pd_df + + +@task +def t2(df: Annotated[pd.DataFrame, my_cols]) -> Annotated[pd.DataFrame, my_cols]: + return df + + +@task +def t3(df: FlyteSchema[my_cols]) -> FlyteSchema[my_cols]: + return df + + +@task +def t4() -> FlyteSchema[my_cols]: + return pd_df + + +@task +def t5(sd: StructuredDataset[my_cols]) -> Annotated[pd.DataFrame, my_cols]: + return sd.open(pd.DataFrame).all() + + +@workflow +def sd_wf() -> Annotated[pd.DataFrame, my_cols]: + df = t1() + return t2(df=df) + + +@workflow +def sd_to_schema_wf() -> pd.DataFrame: + df = t1() + return t3(df=df) + + +@workflow +def schema_to_sd_wf() -> pd.DataFrame: + df = t4() + t2(df=df) + return t5(sd=df) + + +def test_structured_dataset_wf(): + assert_frame_equal(sd_wf(), pd_df) + assert_frame_equal(sd_to_schema_wf(), pd_df) + assert_frame_equal(schema_to_sd_wf(), pd_df) diff --git a/tests/flytekit/unit/models/test_literals.py b/tests/flytekit/unit/models/test_literals.py index b6eec52cd7..45170d9f4d 100644 --- a/tests/flytekit/unit/models/test_literals.py +++ b/tests/flytekit/unit/models/test_literals.py @@ -358,6 +358,43 @@ def test_scalar_schema(): assert len(obj2.value.type.columns) == 6 +def test_structured_dataset(): + my_cols = [ + _types.StructuredDatasetType.DatasetColumn("a", _types.LiteralType(simple=_types.SimpleType.INTEGER)), + _types.StructuredDatasetType.DatasetColumn("b", _types.LiteralType(simple=_types.SimpleType.STRING)), + _types.StructuredDatasetType.DatasetColumn( + "c", _types.LiteralType(collection_type=_types.LiteralType(simple=_types.SimpleType.INTEGER)) + ), + _types.StructuredDatasetType.DatasetColumn( + "d", _types.LiteralType(map_value_type=_types.LiteralType(simple=_types.SimpleType.INTEGER)) + ), + ] + ds = literals.StructuredDataset( + uri="s3://bucket", + metadata=literals.StructuredDatasetMetadata( + structured_dataset_type=_types.StructuredDatasetType(columns=my_cols, format="parquet") + ), + ) + obj = literals.Scalar(structured_dataset=ds) + assert obj.error is None + assert obj.blob is None + assert obj.binary is None + assert obj.schema is None + assert obj.none_type is None + assert obj.structured_dataset is not None + assert obj.value.uri == "s3://bucket" + assert len(obj.value.metadata.structured_dataset_type.columns) == 4 + obj2 = literals.Scalar.from_flyte_idl(obj.to_flyte_idl()) + assert obj == obj2 + assert obj2.blob is None + assert obj2.binary is None + assert obj2.schema is None + assert obj2.none_type is None + assert obj2.structured_dataset is not None + assert obj2.value.uri == "s3://bucket" + assert len(obj2.value.metadata.structured_dataset_type.columns) == 4 + + def test_binding_data_scalar(): obj = literals.BindingData(scalar=literals.Scalar(primitive=literals.Primitive(integer=5))) assert obj.value.value.value == 5 diff --git a/tests/flytekit/unit/type_engines/structured_dataset/test_structured_dataset_workflow.py b/tests/flytekit/unit/type_engines/structured_dataset/test_structured_dataset_workflow.py new file mode 100644 index 0000000000..ae19956b3c --- /dev/null +++ b/tests/flytekit/unit/type_engines/structured_dataset/test_structured_dataset_workflow.py @@ -0,0 +1,224 @@ +import os +import typing + +try: + from typing import Annotated +except ImportError: + from typing_extensions import Annotated + +import numpy as np +import pandas as pd +import pyarrow as pa +import pyarrow.parquet as pq + +from flytekit import FlyteContext, FlyteContextManager, kwtypes, task, workflow +from flytekit.models import literals +from flytekit.models.literals import StructuredDatasetMetadata +from flytekit.models.types import StructuredDatasetType +from flytekit.types.structured.structured_dataset import ( + BIGQUERY, + DF, + FLYTE_DATASET_TRANSFORMER, + LOCAL, + PARQUET, + S3, + StructuredDataset, + StructuredDatasetDecoder, + StructuredDatasetEncoder, +) + +PANDAS_PATH = FlyteContextManager.current_context().file_access.get_random_local_directory() +NUMPY_PATH = FlyteContextManager.current_context().file_access.get_random_local_directory() +BQ_PATH = "bq://flyte-dataset:flyte.table" + +my_cols = kwtypes(w=typing.Dict[str, typing.Dict[str, int]], x=typing.List[typing.List[int]], y=int, z=str) +fields = [("some_int", pa.int32()), ("some_string", pa.string())] +arrow_schema = pa.schema(fields) +pd_df = pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]}) + + +class MockBQEncodingHandlers(StructuredDatasetEncoder): + def encode( + self, + ctx: FlyteContext, + structured_dataset: StructuredDataset, + structured_dataset_type: StructuredDatasetType, + ) -> literals.StructuredDataset: + return literals.StructuredDataset( + uri="s3://bucket/key", metadata=StructuredDatasetMetadata(structured_dataset_type) + ) + + +class MockBQDecodingHandlers(StructuredDatasetDecoder): + def decode( + self, + ctx: FlyteContext, + flyte_value: literals.StructuredDataset, + ) -> pd.DataFrame: + return pd_df + + +FLYTE_DATASET_TRANSFORMER.register_handler(MockBQEncodingHandlers(pd.DataFrame, BIGQUERY), False, True) +FLYTE_DATASET_TRANSFORMER.register_handler(MockBQDecodingHandlers(pd.DataFrame, BIGQUERY), False, True) + + +class NumpyEncodingHandlers(StructuredDatasetEncoder): + def encode( + self, + ctx: FlyteContext, + structured_dataset: StructuredDataset, + structured_dataset_type: StructuredDatasetType, + ) -> literals.StructuredDataset: + path = typing.cast(str, structured_dataset.uri) or ctx.file_access.get_random_remote_directory() + df = typing.cast(np.ndarray, structured_dataset.dataframe) + name = ["col" + str(i) for i in range(len(df))] + table = pa.Table.from_arrays(df, name) + local_dir = ctx.file_access.get_random_local_directory() + local_path = os.path.join(local_dir, f"{0:05}") + pq.write_table(table, local_path) + ctx.file_access.upload_directory(local_dir, path) + structured_dataset_type.format = PARQUET + return literals.StructuredDataset(uri=path, metadata=StructuredDatasetMetadata(structured_dataset_type)) + + +class NumpyDecodingHandlers(StructuredDatasetDecoder): + def decode( + self, + ctx: FlyteContext, + flyte_value: literals.StructuredDataset, + ) -> typing.Union[DF, typing.Generator[DF, None, None]]: + path = flyte_value.uri + local_dir = ctx.file_access.get_random_local_directory() + ctx.file_access.get_data(path, local_dir, is_multipart=True) + table = pq.read_table(local_dir) + return table.to_pandas().to_numpy() + + +for protocol in [LOCAL, S3]: + FLYTE_DATASET_TRANSFORMER.register_handler(NumpyEncodingHandlers(np.ndarray, protocol, PARQUET)) + FLYTE_DATASET_TRANSFORMER.register_handler(NumpyDecodingHandlers(np.ndarray, protocol, PARQUET)) + + +@task +def t1(dataframe: pd.DataFrame) -> Annotated[pd.DataFrame, my_cols]: + # S3 (parquet) -> Pandas -> S3 (parquet) default behaviour + return dataframe + + +@task +def t1a(dataframe: pd.DataFrame) -> StructuredDataset[my_cols, PARQUET]: + # S3 (parquet) -> Pandas -> S3 (parquet) + return StructuredDataset(dataframe=dataframe, uri=PANDAS_PATH) + + +@task +def t2(dataframe: pd.DataFrame) -> Annotated[pd.DataFrame, arrow_schema]: + # S3 (parquet) -> Pandas -> S3 (parquet) + return dataframe + + +@task +def t3(dataset: StructuredDataset[my_cols]) -> StructuredDataset[my_cols]: + # s3 (parquet) -> pandas -> s3 (parquet) + print(dataset.open(pd.DataFrame).all()) + # In the example, we download dataset when we open it. + # Here we won't upload anything, since we're returning just the input object. + return dataset + + +@task +def t3a(dataset: StructuredDataset[my_cols]) -> StructuredDataset[my_cols]: + # This task will not do anything - no uploading, no downloading + return dataset + + +@task +def t4(dataset: StructuredDataset[my_cols]) -> pd.DataFrame: + # s3 (parquet) -> pandas -> s3 (parquet) + return dataset.open(pd.DataFrame).all() + + +@task +def t5(dataframe: pd.DataFrame) -> StructuredDataset[my_cols]: + # s3 (parquet) -> pandas -> bq + return StructuredDataset(dataframe=dataframe, uri=BQ_PATH) + + +@task +def t6(dataset: StructuredDataset[my_cols]) -> pd.DataFrame: + # bq -> pandas -> s3 (parquet) + df = dataset.open(pd.DataFrame).all() + return df + + +@task +def t7(df1: pd.DataFrame, df2: pd.DataFrame) -> (StructuredDataset[my_cols], StructuredDataset[my_cols]): + # df1: pandas -> bq + # df2: pandas -> s3 (parquet) + return StructuredDataset(dataframe=df1, uri=BQ_PATH), StructuredDataset(dataframe=df2) + + +@task +def t8(dataframe: pa.Table) -> StructuredDataset[my_cols]: + # Arrow table -> s3 (parquet) + print(dataframe.columns) + return StructuredDataset(dataframe=dataframe) + + +@task +def t8a(dataframe: pa.Table) -> pa.Table: + # Arrow table -> s3 (parquet) + print(dataframe.columns) + return dataframe + + +@task +def t9(dataframe: np.ndarray) -> StructuredDataset[my_cols]: + # numpy -> Arrow table -> s3 (parquet) + return StructuredDataset(dataframe=dataframe, uri=NUMPY_PATH) + + +@task +def t10(dataset: StructuredDataset[my_cols]) -> np.ndarray: + # s3 (parquet) -> Arrow table -> numpy + np_array = dataset.open(np.ndarray).all() + return np_array + + +@task +def generate_pandas() -> pd.DataFrame: + return pd_df + + +@task +def generate_numpy() -> np.ndarray: + return np.array([[1, 2], [4, 5]]) + + +@task +def generate_arrow() -> pa.Table: + return pa.Table.from_pandas(pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]})) + + +@workflow() +def wf(): + df = generate_pandas() + np_array = generate_numpy() + arrow_df = generate_arrow() + t1(dataframe=df) + t1a(dataframe=df) + t2(dataframe=df) + t3(dataset=StructuredDataset(uri=PANDAS_PATH)) + t3a(dataset=StructuredDataset(uri=PANDAS_PATH)) + t4(dataset=StructuredDataset(uri=PANDAS_PATH)) + t5(dataframe=df) + t6(dataset=StructuredDataset(uri=BQ_PATH)) + t7(df1=df, df2=df) + t8(dataframe=arrow_df) + t8a(dataframe=arrow_df) + t9(dataframe=np_array) + t10(dataset=StructuredDataset(uri=NUMPY_PATH)) + + +def test_structured_dataset_wf(): + wf() From 9dbca8110dce6262d8a03f880431aa2c8d34965e Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Sat, 15 Jan 2022 05:02:34 +0800 Subject: [PATCH 053/128] Parent workflow serialization fails when calling a launch plan with fixed inputs (#814) Signed-off-by: Kevin Su Signed-off-by: maximsmol --- flytekit/common/translator.py | 8 ++++- .../unit/common_tests/test_translator.py | 35 +++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/flytekit/common/translator.py b/flytekit/common/translator.py index d295ef7914..deca82f7fd 100644 --- a/flytekit/common/translator.py +++ b/flytekit/common/translator.py @@ -307,10 +307,16 @@ def get_serializable_node( elif isinstance(entity.flyte_entity, LaunchPlan): lp_spec = get_serializable(entity_mapping, settings, entity.flyte_entity) + # Node's inputs should not contain the data which is fixed input + node_input = [] + for b in entity.bindings: + if b.var not in entity.flyte_entity.fixed_inputs.literals: + node_input.append(b) + node_model = workflow_model.Node( id=_dnsify(entity.id), metadata=entity.metadata, - inputs=entity.bindings, + inputs=node_input, upstream_node_ids=[n.id for n in upstream_sdk_nodes], output_aliases=[], workflow_node=workflow_model.WorkflowNode(launchplan_ref=lp_spec.id), diff --git a/tests/flytekit/unit/common_tests/test_translator.py b/tests/flytekit/unit/common_tests/test_translator.py index b7b048808e..91b1dd2780 100644 --- a/tests/flytekit/unit/common_tests/test_translator.py +++ b/tests/flytekit/unit/common_tests/test_translator.py @@ -124,3 +124,38 @@ def t1(a: int) -> (int, str): ) task_spec = get_serializable(OrderedDict(), ssettings, t2) assert "pyflyte" not in task_spec.template.container.args + + +def test_launch_plan_with_fixed_input(): + @task + def greet(day_of_week: str, number: int, am: bool) -> str: + greeting = "Have a great " + day_of_week + " " + greeting += "morning" if am else "evening" + return greeting + "!" * number + + @workflow + def go_greet(day_of_week: str, number: int, am: bool = False) -> str: + return greet(day_of_week=day_of_week, number=number, am=am) + + morning_greeting = LaunchPlan.create( + "morning_greeting", + go_greet, + fixed_inputs={"am": True}, + default_inputs={"number": 1}, + ) + + @workflow + def morning_greeter_caller(day_of_week: str) -> str: + greeting = morning_greeting(day_of_week=day_of_week) + return greeting + + settings = ( + serialization_settings.new_builder() + .with_fast_serialization_settings(FastSerializationSettings(enabled=True)) + .build() + ) + task_spec = get_serializable(OrderedDict(), settings, morning_greeter_caller) + assert len(task_spec.template.interface.inputs) == 1 + assert len(task_spec.template.interface.outputs) == 1 + assert len(task_spec.template.nodes) == 1 + assert len(task_spec.template.nodes[0].inputs) == 2 From 8a3ca234d6e3fe64bb60e54555c6226e83950966 Mon Sep 17 00:00:00 2001 From: Eduardo Apolinario <653394+eapolinario@users.noreply.github.com> Date: Fri, 14 Jan 2022 15:12:31 -0800 Subject: [PATCH 054/128] Fix sagemaker plugin (#817) Signed-off-by: Eduardo Apolinario Co-authored-by: Eduardo Apolinario Signed-off-by: maximsmol --- plugins/flytekit-aws-sagemaker/setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/plugins/flytekit-aws-sagemaker/setup.py b/plugins/flytekit-aws-sagemaker/setup.py index 1b171e631e..46522b5e9a 100644 --- a/plugins/flytekit-aws-sagemaker/setup.py +++ b/plugins/flytekit-aws-sagemaker/setup.py @@ -16,7 +16,7 @@ author_email="admin@flyte.org", description="AWS Plugins for flytekit", namespace_packages=["flytekitplugins"], - packages=[f"flytekitplugins.{PLUGIN_NAME}"], + packages=[f"flytekitplugins.{PLUGIN_NAME}", f"flytekitplugins.{PLUGIN_NAME}.models"], install_requires=plugin_requires, license="apache2", python_requires=">=3.7", From a63b7ca3715fbd3d1360a1894c16506d3ec66f0a Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Fri, 14 Jan 2022 16:57:36 -0800 Subject: [PATCH 055/128] Remove legacy API (#807) Signed-off-by: Yee Hing Tong --- flytekit/__init__.py | 1 - flytekit/bin/entrypoint.py | 119 +- flytekit/clients/friendly.py | 8 +- flytekit/clients/helpers.py | 4 +- flytekit/clients/raw.py | 9 +- flytekit/clis/flyte_cli/main.py | 283 +--- flytekit/clis/helpers.py | 58 - flytekit/clis/sdk_in_container/basic_auth.py | 6 +- .../clis/sdk_in_container/fast_register.py | 249 ---- flytekit/clis/sdk_in_container/launch_plan.py | 276 ---- flytekit/clis/sdk_in_container/pyflyte.py | 6 - flytekit/clis/sdk_in_container/register.py | 146 -- flytekit/clis/sdk_in_container/serialize.py | 95 +- flytekit/common/component_nodes.py | 157 --- flytekit/common/exceptions/__init__.py | 0 flytekit/common/interface.py | 163 --- flytekit/common/launch_plan.py | 498 ------- flytekit/common/local_workflow.py | 388 ----- flytekit/common/mixins/__init__.py | 0 flytekit/common/mixins/artifact.py | 80 -- flytekit/common/mixins/launchable.py | 129 -- flytekit/common/mixins/registerable.py | 195 --- flytekit/common/nodes.py | 463 ------ flytekit/common/notifications.py | 101 -- flytekit/common/promise.py | 169 --- flytekit/common/schedules.py | 195 --- flytekit/common/sdk_bases.py | 22 - flytekit/common/tasks/__init__.py | 0 flytekit/common/tasks/executions.py | 153 -- flytekit/common/tasks/generic_spark_task.py | 147 -- flytekit/common/tasks/hive_task.py | 299 ---- flytekit/common/tasks/output.py | 46 - flytekit/common/tasks/presto_task.py | 180 --- flytekit/common/tasks/raw_container.py | 237 ---- flytekit/common/tasks/sdk_dynamic.py | 372 ----- flytekit/common/tasks/sdk_runnable.py | 750 ---------- flytekit/common/tasks/sidecar_task.py | 245 ---- flytekit/common/tasks/spark_task.py | 208 --- flytekit/common/tasks/task.py | 423 ------ flytekit/common/types/__init__.py | 3 - flytekit/common/types/base_sdk_types.py | 141 -- flytekit/common/types/blobs.py | 465 ------ flytekit/common/types/containers.py | 157 --- flytekit/common/types/helpers.py | 124 -- flytekit/common/types/impl/__init__.py | 0 flytekit/common/types/impl/blobs.py | 497 ------- flytekit/common/types/impl/schema.py | 995 ------------- flytekit/common/types/primitives.py | 595 -------- flytekit/common/types/proto.py | 319 ----- flytekit/common/types/schema.py | 189 --- flytekit/common/workflow.py | 309 ---- flytekit/common/workflow_execution.py | 183 --- flytekit/configuration/__init__.py | 14 +- flytekit/configuration/common.py | 2 +- flytekit/configuration/platform.py | 2 +- flytekit/contrib/__init__.py | 0 flytekit/contrib/sensors/__init__.py | 0 flytekit/contrib/sensors/base_sensor.py | 90 -- flytekit/contrib/sensors/impl.py | 107 -- flytekit/contrib/sensors/task.py | 127 -- flytekit/core/base_task.py | 9 +- flytekit/{common => core}/constants.py | 0 flytekit/core/container_task.py | 2 +- flytekit/core/context_manager.py | 222 ++- flytekit/core/data_persistence.py | 4 +- flytekit/{common/mixins => core}/hash.py | 0 flytekit/core/interface.py | 2 +- flytekit/core/map_task.py | 4 +- flytekit/{engines/unit => core}/mock_stats.py | 0 flytekit/core/node.py | 2 +- flytekit/core/node_creation.py | 2 +- flytekit/core/promise.py | 4 +- flytekit/core/python_auto_container.py | 2 +- .../core/python_customized_container_task.py | 5 +- flytekit/core/python_function_task.py | 4 +- flytekit/core/reference.py | 2 +- flytekit/core/reference_entity.py | 2 +- flytekit/core/shim_task.py | 3 +- flytekit/core/tracker.py | 2 +- flytekit/core/type_engine.py | 19 +- flytekit/{common => core}/utils.py | 150 +- flytekit/core/workflow.py | 6 +- flytekit/engines/__init__.py | 0 flytekit/engines/common.py | 422 ------ flytekit/engines/flyte/__init__.py | 0 flytekit/engines/flyte/engine.py | 727 ---------- flytekit/engines/loader.py | 44 - flytekit/engines/unit/__init__.py | 0 flytekit/engines/unit/engine.py | 329 ----- flytekit/engines/unit/unit.config | 14 - flytekit/{common => exceptions}/__init__.py | 0 flytekit/{common => }/exceptions/base.py | 0 flytekit/{common => }/exceptions/scopes.py | 6 +- flytekit/{common => }/exceptions/system.py | 2 +- flytekit/{common => }/exceptions/user.py | 4 +- flytekit/extend/__init__.py | 4 +- flytekit/extras/persistence/gcs_gsutil.py | 2 +- flytekit/extras/persistence/http.py | 2 +- flytekit/extras/persistence/s3_awscli.py | 2 +- .../cli_identifiers.py} | 110 +- flytekit/interfaces/data/__init__.py | 0 flytekit/interfaces/data/common.py | 54 - flytekit/interfaces/data/data_proxy.py | 172 --- flytekit/interfaces/data/gcs/__init__.py | 0 flytekit/interfaces/data/gcs/gcs_proxy.py | 150 -- flytekit/interfaces/data/http/__init__.py | 0 .../interfaces/data/http/http_data_proxy.py | 80 -- flytekit/interfaces/data/local/__init__.py | 0 .../interfaces/data/local/local_file_proxy.py | 92 -- flytekit/interfaces/data/s3/__init__.py | 0 flytekit/interfaces/data/s3/s3proxy.py | 224 --- flytekit/models/common.py | 11 +- flytekit/models/core/compiler.py | 9 +- flytekit/models/interface.py | 19 +- flytekit/models/literals.py | 2 +- flytekit/models/task.py | 223 +-- flytekit/plugins/__init__.py | 34 - flytekit/remote/component_nodes.py | 2 +- flytekit/remote/executions.py | 4 +- flytekit/remote/launch_plan.py | 12 - flytekit/remote/nodes.py | 8 +- flytekit/remote/remote.py | 14 +- flytekit/remote/task.py | 2 +- flytekit/remote/workflow.py | 6 +- flytekit/sdk/__init__.py | 0 flytekit/sdk/exceptions.py | 11 - flytekit/sdk/spark_types.py | 8 - flytekit/sdk/tasks.py | 1244 ----------------- flytekit/sdk/test_utils.py | 92 -- flytekit/sdk/types.py | 504 ------- flytekit/sdk/workflow.py | 132 -- flytekit/testing/__init__.py | 1 - flytekit/tools/fast_registration.py | 11 +- flytekit/tools/lazy_loader.py | 118 -- flytekit/tools/module_loader.py | 136 +- flytekit/{common => tools}/translator.py | 4 +- flytekit/type_engines/__init__.py | 0 flytekit/type_engines/common.py | 31 - flytekit/type_engines/default/__init__.py | 0 flytekit/type_engines/default/flyte.py | 193 --- flytekit/types/schema/types.py | 2 +- .../flytekitplugins/awssagemaker/hpo.py | 7 +- .../awssagemaker/models/parameter_ranges.py | 2 +- .../flytekit-aws-sagemaker/tests/test_hpo.py | 6 +- .../tests/test_training.py | 2 +- .../flytekitplugins/pod/task.py | 2 +- plugins/flytekit-k8s-pod/tests/test_pod.py | 2 +- .../flytekitplugins/papermill/task.py | 2 +- .../flytekitplugins/spark/models.py | 27 +- .../flytekitplugins/spark/task.py | 5 +- .../flytekit-spark/tests/test_spark_task.py | 2 +- .../tests/test_sql_tracker.py | 2 +- .../flytekit-sqlalchemy/tests/test_task.py | 2 +- setup.py | 38 +- tests/flytekit/common/parameterizers.py | 36 +- tests/flytekit/common/task_definitions.py | 9 - tests/flytekit/common/workflows/__init__.py | 0 tests/flytekit/common/workflows/batch.py | 128 -- .../common/workflows/dynamic_workflows.py | 39 - .../common/workflows/failing_workflows.py | 28 - tests/flytekit/common/workflows/gpu.py | 20 - tests/flytekit/common/workflows/hive.py | 33 - tests/flytekit/common/workflows/nested.py | 49 - tests/flytekit/common/workflows/notebook.py | 25 - .../common/workflows/notifications.py | 34 - tests/flytekit/common/workflows/presto.py | 25 - tests/flytekit/common/workflows/python.py | 67 - .../common/workflows/raw_container.py | 31 - .../common/workflows/raw_edge_detector.py | 20 - .../flytekit/common/workflows/scala_spark.py | 32 - tests/flytekit/common/workflows/sidecar.py | 56 - tests/flytekit/common/workflows/simple.py | 114 -- tests/flytekit/common/workflows/spark.py | 50 - .../integration/remote/test_remote.py | 2 +- tests/flytekit/loadtests/__init__.py | 0 tests/flytekit/loadtests/cp_orchestrator.py | 29 - tests/flytekit/loadtests/cp_python.py | 28 - tests/flytekit/loadtests/cp_spark.py | 42 - tests/flytekit/loadtests/dynamic_job.py | 40 - tests/flytekit/loadtests/orchestrator.py | 48 - tests/flytekit/loadtests/python.py | 27 - .../unit/bin/test_python_entrypoint.py | 204 +-- .../unit/cli/pyflyte/test_launch_plans.py | 30 - .../flytekit/unit/cli/pyflyte/test_package.py | 2 +- .../unit/cli/pyflyte/test_register.py | 40 - tests/flytekit/unit/cli/test_cli_helpers.py | 33 - tests/flytekit/unit/cli/test_flyte_cli.py | 43 +- tests/flytekit/unit/common_tests/__init__.py | 0 .../unit/common_tests/exceptions/__init__.py | 0 .../unit/common_tests/mixins/__init__.py | 0 .../mixins/sample_registerable.py | 15 - .../common_tests/mixins/test_registerable.py | 13 - .../unit/common_tests/tasks/__init__.py | 0 .../unit/common_tests/tasks/spark/__init__.py | 0 .../tasks/spark/test_spark_task.py | 37 - .../tasks/test_execution_params.py | 76 - .../tasks/test_raw_container_task.py | 27 - .../common_tests/tasks/test_sdk_runnable.py | 44 - .../unit/common_tests/tasks/test_task.py | 108 -- .../unit/common_tests/test_interface.py | 80 -- .../unit/common_tests/test_launch_plan.py | 388 ----- .../flytekit/unit/common_tests/test_nodes.py | 292 ---- .../unit/common_tests/test_notifications.py | 35 - .../unit/common_tests/test_promise.py | 89 -- .../unit/common_tests/test_schedules.py | 131 -- .../unit/common_tests/test_workflow.py | 383 ----- .../common_tests/test_workflow_promote.py | 101 -- .../unit/common_tests/types/impl/__init__.py | 0 .../common_tests/types/impl/test_blobs.py | 332 ----- .../common_tests/types/impl/test_schema.py | 553 -------- .../unit/common_tests/types/test_blobs.py | 141 -- .../common_tests/types/test_containers.py | 173 --- .../unit/common_tests/types/test_helpers.py | 59 - .../common_tests/types/test_primitives.py | 289 ---- .../unit/common_tests/types/test_proto.py | 63 - .../unit/common_tests/types/test_schema.py | 69 - .../unit/configuration/test_waterfall.py | 2 +- tests/flytekit/unit/contrib/__init__.py | 0 .../flytekit/unit/contrib/sensors/__init__.py | 0 .../unit/contrib/sensors/test_impl.py | 58 - .../unit/contrib/sensors/test_task.py | 22 - tests/flytekit/unit/core/test_conditions.py | 2 +- .../unit/core/test_flyte_directory.py | 2 +- tests/flytekit/unit/core/test_flyte_pickle.py | 2 +- tests/flytekit/unit/core/test_imperative.py | 4 +- tests/flytekit/unit/core/test_launch_plan.py | 2 +- tests/flytekit/unit/core/test_map_task.py | 2 +- .../flytekit/unit/core/test_node_creation.py | 4 +- tests/flytekit/unit/core/test_references.py | 2 +- tests/flytekit/unit/core/test_resolver.py | 2 +- .../flytekit/unit/core/test_serialization.py | 2 +- tests/flytekit/unit/core/test_type_engine.py | 1 + tests/flytekit/unit/core/test_type_hints.py | 1 + .../unit/{common_tests => core}/test_utils.py | 2 +- tests/flytekit/unit/core/test_workflows.py | 4 +- tests/flytekit/unit/engines/__init__.py | 0 tests/flytekit/unit/engines/flyte/__init__.py | 0 .../unit/engines/flyte/test_engine.py | 807 ----------- tests/flytekit/unit/engines/test_loader.py | 13 - tests/flytekit/unit/engines/unit/__init__.py | 0 .../flytekit/unit/exceptions}/__init__.py | 0 .../exceptions/test_base.py | 2 +- .../exceptions/test_scopes.py | 2 +- .../exceptions/test_system.py | 2 +- .../exceptions/test_user.py | 2 +- .../unit/extras/sqlite3/test_sql_tracker.py | 2 +- .../interfaces/data/gcs/test_gcs_proxy.py | 80 -- .../unit/interfaces/data/s3/test_s3_proxy.py | 36 - .../unit/models/test_dynamic_spark.py | 29 - .../flytekit/unit/models/test_dynamic_wfs.py | 125 -- tests/flytekit/unit/models/test_tasks.py | 31 - tests/flytekit/unit/remote/test_remote.py | 2 +- .../unit/remote/test_wrapper_classes.py | 2 +- tests/flytekit/unit/sdk/__init__.py | 0 tests/flytekit/unit/sdk/conftest.py | 9 - tests/flytekit/unit/sdk/tasks/__init__.py | 0 .../sdk/tasks/test_dynamic_sidecar_tasks.py | 82 -- .../unit/sdk/tasks/test_dynamic_tasks.py | 241 ---- .../unit/sdk/tasks/test_hive_tasks.py | 141 -- .../unit/sdk/tasks/test_sidecar_tasks.py | 84 -- .../unit/sdk/tasks/test_spark_task.py | 78 -- tests/flytekit/unit/sdk/tasks/test_tasks.py | 108 -- tests/flytekit/unit/sdk/test_workflow.py | 150 -- tests/flytekit/unit/sdk/types/__init__.py | 0 tests/flytekit/unit/sdk/types/test_blobs.py | 31 - .../unit/sdk/types/test_primitives.py | 33 - tests/flytekit/unit/sdk/types/test_schema.py | 49 - tests/flytekit/unit/tasks/__init__.py | 0 tests/flytekit/unit/test_plugins.py | 48 - .../{common_tests => }/test_translator.py | 2 +- tests/flytekit/unit/tools/test_aws.py | 7 - tests/flytekit/unit/tools/test_lazy_loader.py | 14 - .../flytekit/unit/tools/test_module_loader.py | 4 +- tests/flytekit/unit/type_engines/__init__.py | 0 .../unit/type_engines/default/__init__.py | 0 .../default/test_flyte_type_engine.py | 67 - tests/flytekit/unit/use_scenarios/__init__.py | 0 .../use_scenarios/unit_testing/__init__.py | 0 .../use_scenarios/unit_testing/test_blobs.py | 165 --- .../unit_testing/test_hive_tasks.py | 44 - .../unit_testing/test_schemas.py | 139 -- 281 files changed, 656 insertions(+), 23968 deletions(-) delete mode 100644 flytekit/clis/sdk_in_container/fast_register.py delete mode 100644 flytekit/clis/sdk_in_container/launch_plan.py delete mode 100644 flytekit/clis/sdk_in_container/register.py delete mode 100644 flytekit/common/component_nodes.py delete mode 100644 flytekit/common/exceptions/__init__.py delete mode 100644 flytekit/common/interface.py delete mode 100644 flytekit/common/launch_plan.py delete mode 100644 flytekit/common/local_workflow.py delete mode 100644 flytekit/common/mixins/__init__.py delete mode 100644 flytekit/common/mixins/artifact.py delete mode 100644 flytekit/common/mixins/launchable.py delete mode 100644 flytekit/common/mixins/registerable.py delete mode 100644 flytekit/common/nodes.py delete mode 100644 flytekit/common/notifications.py delete mode 100644 flytekit/common/promise.py delete mode 100644 flytekit/common/schedules.py delete mode 100644 flytekit/common/sdk_bases.py delete mode 100644 flytekit/common/tasks/__init__.py delete mode 100644 flytekit/common/tasks/executions.py delete mode 100644 flytekit/common/tasks/generic_spark_task.py delete mode 100644 flytekit/common/tasks/hive_task.py delete mode 100644 flytekit/common/tasks/output.py delete mode 100644 flytekit/common/tasks/presto_task.py delete mode 100644 flytekit/common/tasks/raw_container.py delete mode 100644 flytekit/common/tasks/sdk_dynamic.py delete mode 100644 flytekit/common/tasks/sdk_runnable.py delete mode 100644 flytekit/common/tasks/sidecar_task.py delete mode 100644 flytekit/common/tasks/spark_task.py delete mode 100644 flytekit/common/tasks/task.py delete mode 100644 flytekit/common/types/__init__.py delete mode 100644 flytekit/common/types/base_sdk_types.py delete mode 100644 flytekit/common/types/blobs.py delete mode 100644 flytekit/common/types/containers.py delete mode 100644 flytekit/common/types/helpers.py delete mode 100644 flytekit/common/types/impl/__init__.py delete mode 100644 flytekit/common/types/impl/blobs.py delete mode 100644 flytekit/common/types/impl/schema.py delete mode 100644 flytekit/common/types/primitives.py delete mode 100644 flytekit/common/types/proto.py delete mode 100644 flytekit/common/types/schema.py delete mode 100644 flytekit/common/workflow.py delete mode 100644 flytekit/common/workflow_execution.py delete mode 100644 flytekit/contrib/__init__.py delete mode 100644 flytekit/contrib/sensors/__init__.py delete mode 100644 flytekit/contrib/sensors/base_sensor.py delete mode 100644 flytekit/contrib/sensors/impl.py delete mode 100644 flytekit/contrib/sensors/task.py rename flytekit/{common => core}/constants.py (100%) rename flytekit/{common/mixins => core}/hash.py (100%) rename flytekit/{engines/unit => core}/mock_stats.py (100%) rename flytekit/{common => core}/utils.py (57%) delete mode 100644 flytekit/engines/__init__.py delete mode 100644 flytekit/engines/common.py delete mode 100644 flytekit/engines/flyte/__init__.py delete mode 100644 flytekit/engines/flyte/engine.py delete mode 100644 flytekit/engines/loader.py delete mode 100644 flytekit/engines/unit/__init__.py delete mode 100644 flytekit/engines/unit/engine.py delete mode 100644 flytekit/engines/unit/unit.config rename flytekit/{common => exceptions}/__init__.py (100%) rename flytekit/{common => }/exceptions/base.py (100%) rename flytekit/{common => }/exceptions/scopes.py (97%) rename flytekit/{common => }/exceptions/system.py (95%) rename flytekit/{common => }/exceptions/user.py (94%) rename flytekit/{common/core/identifier.py => interfaces/cli_identifiers.py} (89%) delete mode 100644 flytekit/interfaces/data/__init__.py delete mode 100644 flytekit/interfaces/data/common.py delete mode 100644 flytekit/interfaces/data/data_proxy.py delete mode 100644 flytekit/interfaces/data/gcs/__init__.py delete mode 100644 flytekit/interfaces/data/gcs/gcs_proxy.py delete mode 100644 flytekit/interfaces/data/http/__init__.py delete mode 100644 flytekit/interfaces/data/http/http_data_proxy.py delete mode 100644 flytekit/interfaces/data/local/__init__.py delete mode 100644 flytekit/interfaces/data/local/local_file_proxy.py delete mode 100644 flytekit/interfaces/data/s3/__init__.py delete mode 100644 flytekit/interfaces/data/s3/s3proxy.py delete mode 100644 flytekit/plugins/__init__.py delete mode 100644 flytekit/sdk/__init__.py delete mode 100644 flytekit/sdk/exceptions.py delete mode 100644 flytekit/sdk/spark_types.py delete mode 100644 flytekit/sdk/tasks.py delete mode 100644 flytekit/sdk/test_utils.py delete mode 100644 flytekit/sdk/types.py delete mode 100644 flytekit/sdk/workflow.py delete mode 100644 flytekit/tools/lazy_loader.py rename flytekit/{common => tools}/translator.py (99%) delete mode 100644 flytekit/type_engines/__init__.py delete mode 100644 flytekit/type_engines/common.py delete mode 100644 flytekit/type_engines/default/__init__.py delete mode 100644 flytekit/type_engines/default/flyte.py delete mode 100644 tests/flytekit/common/task_definitions.py delete mode 100644 tests/flytekit/common/workflows/__init__.py delete mode 100644 tests/flytekit/common/workflows/batch.py delete mode 100644 tests/flytekit/common/workflows/dynamic_workflows.py delete mode 100644 tests/flytekit/common/workflows/failing_workflows.py delete mode 100644 tests/flytekit/common/workflows/gpu.py delete mode 100644 tests/flytekit/common/workflows/hive.py delete mode 100644 tests/flytekit/common/workflows/nested.py delete mode 100644 tests/flytekit/common/workflows/notebook.py delete mode 100644 tests/flytekit/common/workflows/notifications.py delete mode 100644 tests/flytekit/common/workflows/presto.py delete mode 100644 tests/flytekit/common/workflows/python.py delete mode 100644 tests/flytekit/common/workflows/raw_container.py delete mode 100644 tests/flytekit/common/workflows/raw_edge_detector.py delete mode 100644 tests/flytekit/common/workflows/scala_spark.py delete mode 100644 tests/flytekit/common/workflows/sidecar.py delete mode 100644 tests/flytekit/common/workflows/simple.py delete mode 100644 tests/flytekit/common/workflows/spark.py delete mode 100644 tests/flytekit/loadtests/__init__.py delete mode 100644 tests/flytekit/loadtests/cp_orchestrator.py delete mode 100644 tests/flytekit/loadtests/cp_python.py delete mode 100644 tests/flytekit/loadtests/cp_spark.py delete mode 100644 tests/flytekit/loadtests/dynamic_job.py delete mode 100644 tests/flytekit/loadtests/orchestrator.py delete mode 100644 tests/flytekit/loadtests/python.py delete mode 100644 tests/flytekit/unit/cli/pyflyte/test_launch_plans.py delete mode 100644 tests/flytekit/unit/cli/pyflyte/test_register.py delete mode 100644 tests/flytekit/unit/common_tests/__init__.py delete mode 100644 tests/flytekit/unit/common_tests/exceptions/__init__.py delete mode 100644 tests/flytekit/unit/common_tests/mixins/__init__.py delete mode 100644 tests/flytekit/unit/common_tests/mixins/sample_registerable.py delete mode 100644 tests/flytekit/unit/common_tests/mixins/test_registerable.py delete mode 100644 tests/flytekit/unit/common_tests/tasks/__init__.py delete mode 100644 tests/flytekit/unit/common_tests/tasks/spark/__init__.py delete mode 100644 tests/flytekit/unit/common_tests/tasks/spark/test_spark_task.py delete mode 100644 tests/flytekit/unit/common_tests/tasks/test_execution_params.py delete mode 100644 tests/flytekit/unit/common_tests/tasks/test_raw_container_task.py delete mode 100644 tests/flytekit/unit/common_tests/tasks/test_sdk_runnable.py delete mode 100644 tests/flytekit/unit/common_tests/tasks/test_task.py delete mode 100644 tests/flytekit/unit/common_tests/test_interface.py delete mode 100644 tests/flytekit/unit/common_tests/test_launch_plan.py delete mode 100644 tests/flytekit/unit/common_tests/test_nodes.py delete mode 100644 tests/flytekit/unit/common_tests/test_notifications.py delete mode 100644 tests/flytekit/unit/common_tests/test_promise.py delete mode 100644 tests/flytekit/unit/common_tests/test_schedules.py delete mode 100644 tests/flytekit/unit/common_tests/test_workflow.py delete mode 100644 tests/flytekit/unit/common_tests/types/impl/__init__.py delete mode 100644 tests/flytekit/unit/common_tests/types/impl/test_blobs.py delete mode 100644 tests/flytekit/unit/common_tests/types/impl/test_schema.py delete mode 100644 tests/flytekit/unit/common_tests/types/test_blobs.py delete mode 100644 tests/flytekit/unit/common_tests/types/test_containers.py delete mode 100644 tests/flytekit/unit/common_tests/types/test_helpers.py delete mode 100644 tests/flytekit/unit/common_tests/types/test_primitives.py delete mode 100644 tests/flytekit/unit/common_tests/types/test_proto.py delete mode 100644 tests/flytekit/unit/common_tests/types/test_schema.py delete mode 100644 tests/flytekit/unit/contrib/__init__.py delete mode 100644 tests/flytekit/unit/contrib/sensors/__init__.py delete mode 100644 tests/flytekit/unit/contrib/sensors/test_impl.py delete mode 100644 tests/flytekit/unit/contrib/sensors/test_task.py rename tests/flytekit/unit/{common_tests => core}/test_utils.py (91%) delete mode 100644 tests/flytekit/unit/engines/__init__.py delete mode 100644 tests/flytekit/unit/engines/flyte/__init__.py delete mode 100644 tests/flytekit/unit/engines/flyte/test_engine.py delete mode 100644 tests/flytekit/unit/engines/test_loader.py delete mode 100644 tests/flytekit/unit/engines/unit/__init__.py rename {flytekit/common/core => tests/flytekit/unit/exceptions}/__init__.py (100%) rename tests/flytekit/unit/{common_tests => }/exceptions/test_base.py (85%) rename tests/flytekit/unit/{common_tests => }/exceptions/test_scopes.py (98%) rename tests/flytekit/unit/{common_tests => }/exceptions/test_system.py (98%) rename tests/flytekit/unit/{common_tests => }/exceptions/test_user.py (98%) delete mode 100644 tests/flytekit/unit/interfaces/data/gcs/test_gcs_proxy.py delete mode 100644 tests/flytekit/unit/interfaces/data/s3/test_s3_proxy.py delete mode 100644 tests/flytekit/unit/models/test_dynamic_spark.py delete mode 100644 tests/flytekit/unit/models/test_dynamic_wfs.py delete mode 100644 tests/flytekit/unit/sdk/__init__.py delete mode 100644 tests/flytekit/unit/sdk/conftest.py delete mode 100644 tests/flytekit/unit/sdk/tasks/__init__.py delete mode 100644 tests/flytekit/unit/sdk/tasks/test_dynamic_sidecar_tasks.py delete mode 100644 tests/flytekit/unit/sdk/tasks/test_dynamic_tasks.py delete mode 100644 tests/flytekit/unit/sdk/tasks/test_hive_tasks.py delete mode 100644 tests/flytekit/unit/sdk/tasks/test_sidecar_tasks.py delete mode 100644 tests/flytekit/unit/sdk/tasks/test_spark_task.py delete mode 100644 tests/flytekit/unit/sdk/tasks/test_tasks.py delete mode 100644 tests/flytekit/unit/sdk/test_workflow.py delete mode 100644 tests/flytekit/unit/sdk/types/__init__.py delete mode 100644 tests/flytekit/unit/sdk/types/test_blobs.py delete mode 100644 tests/flytekit/unit/sdk/types/test_primitives.py delete mode 100644 tests/flytekit/unit/sdk/types/test_schema.py delete mode 100644 tests/flytekit/unit/tasks/__init__.py delete mode 100644 tests/flytekit/unit/test_plugins.py rename tests/flytekit/unit/{common_tests => }/test_translator.py (99%) delete mode 100644 tests/flytekit/unit/tools/test_aws.py delete mode 100644 tests/flytekit/unit/tools/test_lazy_loader.py delete mode 100644 tests/flytekit/unit/type_engines/__init__.py delete mode 100644 tests/flytekit/unit/type_engines/default/__init__.py delete mode 100644 tests/flytekit/unit/type_engines/default/test_flyte_type_engine.py delete mode 100644 tests/flytekit/unit/use_scenarios/__init__.py delete mode 100644 tests/flytekit/unit/use_scenarios/unit_testing/__init__.py delete mode 100644 tests/flytekit/unit/use_scenarios/unit_testing/test_blobs.py delete mode 100644 tests/flytekit/unit/use_scenarios/unit_testing/test_hive_tasks.py delete mode 100644 tests/flytekit/unit/use_scenarios/unit_testing/test_schemas.py diff --git a/flytekit/__init__.py b/flytekit/__init__.py index 3a165084b0..369316c642 100644 --- a/flytekit/__init__.py +++ b/flytekit/__init__.py @@ -160,7 +160,6 @@ else: from importlib.metadata import entry_points -import flytekit.plugins # This will be deprecated, these are the old plugins, the new plugins live in plugins/ from flytekit.core.base_sql_task import SQLTask from flytekit.core.base_task import SecurityContext, TaskMetadata, kwtypes from flytekit.core.condition import conditional diff --git a/flytekit/bin/entrypoint.py b/flytekit/bin/entrypoint.py index 1888376e50..754eab666d 100644 --- a/flytekit/bin/entrypoint.py +++ b/flytekit/bin/entrypoint.py @@ -1,10 +1,8 @@ import contextlib import datetime as _datetime -import importlib as _importlib import logging as python_logging import os as _os import pathlib -import random as _random import traceback as _traceback from typing import List @@ -12,18 +10,14 @@ from flyteidl.core import literals_pb2 as _literals_pb2 from flytekit import PythonFunctionTask -from flytekit.common import constants as _constants -from flytekit.common import utils as _common_utils -from flytekit.common import utils as _utils -from flytekit.common.exceptions import scopes as _scoped_exceptions -from flytekit.common.exceptions import scopes as _scopes -from flytekit.common.exceptions import system as _system_exceptions -from flytekit.common.tasks.sdk_runnable import ExecutionParameters from flytekit.configuration import TemporaryConfiguration as _TemporaryConfiguration from flytekit.configuration import internal as _internal_config from flytekit.configuration import sdk as _sdk_config +from flytekit.core import constants as _constants +from flytekit.core import utils from flytekit.core.base_task import IgnoreOutputs, PythonTask from flytekit.core.context_manager import ( + ExecutionParameters, ExecutionState, FlyteContext, FlyteContextManager, @@ -33,9 +27,8 @@ from flytekit.core.data_persistence import FileAccessProvider from flytekit.core.map_task import MapPythonTask from flytekit.core.promise import VoidPromise -from flytekit.engines import loader as _engine_loader -from flytekit.interfaces import random as _flyte_random -from flytekit.interfaces.data import data_proxy as _data_proxy +from flytekit.exceptions import scopes as _scoped_exceptions +from flytekit.exceptions import scopes as _scopes from flytekit.interfaces.stats.taggable import get_stats as _get_stats from flytekit.loggers import entrypoint_logger as logger from flytekit.models import dynamic_job as _dynamic_job @@ -47,6 +40,12 @@ from flytekit.tools.module_loader import load_object_from_module +def get_version_message(): + import flytekit + + return f"Welcome to Flyte! Version: {flytekit.__version__}" + + def _compute_array_job_index(): # type () -> int """ @@ -61,25 +60,6 @@ def _compute_array_job_index(): return offset + int(_os.environ.get(_os.environ.get("BATCH_JOB_ARRAY_INDEX_VAR_NAME"))) -def _map_job_index_to_child_index(local_input_dir, datadir, index): - local_lookup_file = local_input_dir.get_named_tempfile("indexlookup.pb") - idx_lookup_file = _os.path.join(datadir, "indexlookup.pb") - - # if the indexlookup.pb does not exist, then just return the index - if not _data_proxy.Data.data_exists(idx_lookup_file): - return index - - _data_proxy.Data.get_data(idx_lookup_file, local_lookup_file) - mapping_proto = _utils.load_proto_from_file(_literals_pb2.LiteralCollection, local_lookup_file) - if len(mapping_proto.literals) < index: - raise _system_exceptions.FlyteSystemAssertion( - "dynamic task index lookup array size: {} is smaller than lookup index {}".format( - len(mapping_proto.literals), index - ) - ) - return mapping_proto.literals[index].scalar.primitive.integer - - def _dispatch_execute( ctx: FlyteContext, task_def: PythonTask, @@ -101,7 +81,7 @@ def _dispatch_execute( # Step1 local_inputs_file = _os.path.join(ctx.execution_state.working_dir, "inputs.pb") ctx.file_access.get_data(inputs_path, local_inputs_file) - input_proto = _utils.load_proto_from_file(_literals_pb2.LiteralMap, local_inputs_file) + input_proto = utils.load_proto_from_file(_literals_pb2.LiteralMap, local_inputs_file) idl_input_literals = _literal_models.LiteralMap.from_flyte_idl(input_proto) # Step2 @@ -174,7 +154,7 @@ def _dispatch_execute( logger.error("!! End Error Captured by Flyte !!") for k, v in output_file_dict.items(): - _common_utils.write_proto_to_file(v.to_flyte_idl(), _os.path.join(ctx.execution_state.engine_dir, k)) + utils.write_proto_to_file(v.to_flyte_idl(), _os.path.join(ctx.execution_state.engine_dir, k)) ctx.file_access.put_data(ctx.execution_state.engine_dir, output_prefix, is_multipart=True) logger.info(f"Engine folder written successfully to the output prefix {output_prefix}") @@ -283,45 +263,6 @@ def _handle_annotated_task( _dispatch_execute(ctx, task_def, inputs, output_prefix) -@_scopes.system_entry_point -def _legacy_execute_task(task_module, task_name, inputs, output_prefix, raw_output_data_prefix, test): - """ - This function should be called for old flytekit api tasks (the only API that was available in 0.15.x and earlier) - """ - with _TemporaryConfiguration(_internal_config.CONFIGURATION_PATH.get()): - with _utils.AutoDeletingTempDir("input_dir") as input_dir: - # Load user code - task_module = _importlib.import_module(task_module) - task_def = getattr(task_module, task_name) - - local_inputs_file = input_dir.get_named_tempfile("inputs.pb") - - # Handle inputs/outputs for array job. - if _os.environ.get("BATCH_JOB_ARRAY_INDEX_VAR_NAME"): - job_index = _compute_array_job_index() - - # TODO: Perhaps remove. This is a workaround to an issue we perceived with limited entropy in - # TODO: AWS batch array jobs. - _flyte_random.seed_flyte_random( - "{} {} {}".format(_random.random(), _datetime.datetime.utcnow(), job_index) - ) - - # If an ArrayTask is discoverable, the original job index may be different than the one specified in - # the environment variable. Look up the correct input/outputs in the index lookup mapping file. - job_index = _map_job_index_to_child_index(input_dir, inputs, job_index) - - inputs = _os.path.join(inputs, str(job_index), "inputs.pb") - output_prefix = _os.path.join(output_prefix, str(job_index)) - - _data_proxy.Data.get_data(inputs, local_inputs_file) - input_proto = _utils.load_proto_from_file(_literals_pb2.LiteralMap, local_inputs_file) - - _engine_loader.get_engine().get_task(task_def).execute( - _literal_models.LiteralMap.from_flyte_idl(input_proto), - context={"output_prefix": output_prefix, "raw_output_data_prefix": raw_output_data_prefix}, - ) - - @_scopes.system_entry_point def _execute_task( inputs, @@ -416,8 +357,6 @@ def _pass_through(): @_pass_through.command("pyflyte-execute") -@_click.option("--task-module", required=False) -@_click.option("--task-name", required=False) @_click.option("--inputs", required=True) @_click.option("--output-prefix", required=True) @_click.option("--raw-output-data-prefix", required=False) @@ -431,8 +370,6 @@ def _pass_through(): nargs=-1, ) def execute_task_cmd( - task_module, - task_name, inputs, output_prefix, raw_output_data_prefix, @@ -442,7 +379,7 @@ def execute_task_cmd( resolver, resolver_args, ): - logger.info(_utils.get_version_message()) + logger.info(get_version_message()) # We get weird errors if there are no click echo messages at all, so emit an empty string so that unit tests pass. _click.echo("") # Backwards compatibility - if Propeller hasn't filled this in, then it'll come through here as the original @@ -455,21 +392,17 @@ def execute_task_cmd( # Use the presence of the resolver to differentiate between old API tasks and new API tasks # The addition of a new top-level command seemed out of scope at the time of this writing to pursue given how # pervasive this top level command already (plugins mostly). - if not resolver: - logger.info("No resolver found, assuming legacy API task...") - _legacy_execute_task(task_module, task_name, inputs, output_prefix, raw_output_data_prefix, test) - else: - logger.debug(f"Running task execution with resolver {resolver}...") - _execute_task( - inputs, - output_prefix, - raw_output_data_prefix, - test, - resolver, - resolver_args, - dynamic_addl_distro, - dynamic_dest_dir, - ) + logger.debug(f"Running task execution with resolver {resolver}...") + _execute_task( + inputs, + output_prefix, + raw_output_data_prefix, + test, + resolver, + resolver_args, + dynamic_addl_distro, + dynamic_dest_dir, + ) @_pass_through.command("pyflyte-fast-execute") @@ -528,7 +461,7 @@ def map_execute_task_cmd( resolver, resolver_args, ): - logger.info(_utils.get_version_message()) + logger.info(get_version_message()) _execute_map_task( inputs, diff --git a/flytekit/clients/friendly.py b/flytekit/clients/friendly.py index 5a6bbc4fdc..5daa009ad5 100644 --- a/flytekit/clients/friendly.py +++ b/flytekit/clients/friendly.py @@ -560,7 +560,7 @@ def create_execution(self, project, domain, name, execution_spec, inputs): def recover_execution(self, id, name: str = None): """ Recreates a previously-run workflow execution that will only start executing from the last known failure point. - :param flytekit.common.core.identifier.WorkflowExecutionIdentifier id: + :param flytekit.models.core.identifier.WorkflowExecutionIdentifier id: :param name str: Optional name to assign to the newly created execution. :rtype: flytekit.models.core.identifier.WorkflowExecutionIdentifier """ @@ -572,7 +572,7 @@ def recover_execution(self, id, name: str = None): def get_execution(self, id): """ - :param flytekit.common.core.identifier.WorkflowExecutionIdentifier id: + :param flytekit.models.core.identifier.WorkflowExecutionIdentifier id: :rtype: flytekit.models.execution.Execution """ return _execution.Execution.from_flyte_idl( @@ -638,7 +638,7 @@ def list_executions_paginated(self, project, domain, limit=100, token=None, filt def terminate_execution(self, id, cause): """ - :param flytekit.common.core.identifier.WorkflowExecutionIdentifier id: + :param flytekit.models.core.identifier.WorkflowExecutionIdentifier id: :param Text cause: """ super(SynchronousFlyteClient, self).terminate_execution( @@ -647,7 +647,7 @@ def terminate_execution(self, id, cause): def relaunch_execution(self, id, name=None): """ - :param flytekit.common.core.identifier.WorkflowExecutionIdentifier id: + :param flytekit.models.core.identifier.WorkflowExecutionIdentifier id: :param Text name: [Optional] name for the new execution. If not specified, a randomly generated name will be used :returns: The unique identifier for the new execution. diff --git a/flytekit/clients/helpers.py b/flytekit/clients/helpers.py index 4e2dc71e25..2df64f080e 100644 --- a/flytekit/clients/helpers.py +++ b/flytekit/clients/helpers.py @@ -9,8 +9,8 @@ def iterate_node_executions( """ This returns a generator for node executions. :param flytekit.clients.friendly.SynchronousFlyteClient client: - :param flytekit.common.core.identifier.WorkflowExecutionIdentifier workflow_execution_identifier: - :param flytekit.common.core.identifier.TaskExecutionIdentifier task_execution_identifier: + :param flytekit.models.core.identifier.WorkflowExecutionIdentifier workflow_execution_identifier: + :param flytekit.models.core.identifier.TaskExecutionIdentifier task_execution_identifier: :param int limit: The maximum number of elements to retrieve :param list[flytekit.models.filters.Filter] filters: :rtype: Iterator[flytekit.models.node_execution.NodeExecution] diff --git a/flytekit/clients/raw.py b/flytekit/clients/raw.py index b8f7f578cf..ce58d8cc57 100644 --- a/flytekit/clients/raw.py +++ b/flytekit/clients/raw.py @@ -2,7 +2,6 @@ import time from typing import List -import six as _six from flyteidl.service import admin_pb2_grpc as _admin_service from google.protobuf.json_format import MessageToJson as _MessageToJson from grpc import RpcError as _RpcError @@ -13,13 +12,13 @@ from flytekit.clis.auth import credentials as _credentials_access from flytekit.clis.sdk_in_container import basic_auth as _basic_auth -from flytekit.common.exceptions import user as _user_exceptions from flytekit.configuration import creds as _creds_config from flytekit.configuration.creds import _DEPRECATED_CLIENT_CREDENTIALS_SCOPE as _DEPRECATED_SCOPE from flytekit.configuration.creds import CLIENT_ID as _CLIENT_ID from flytekit.configuration.creds import COMMAND as _COMMAND from flytekit.configuration.creds import DEPRECATED_OAUTH_SCOPES, SCOPES from flytekit.configuration.platform import AUTH as _AUTH +from flytekit.exceptions import user as _user_exceptions from flytekit.loggers import cli_logger @@ -91,9 +90,7 @@ def _refresh_credentials_from_command(flyte_client): output = subprocess.run(command, capture_output=True, text=True, check=True) except subprocess.CalledProcessError as e: cli_logger.error("Failed to generate token from command {}".format(command)) - raise _user_exceptions.FlyteAuthenticationException( - "Problems refreshing token with command: " + _six.text_type(e) - ) + raise _user_exceptions.FlyteAuthenticationException("Problems refreshing token with command: " + str(e)) flyte_client.set_access_token(output.stdout.strip()) @@ -134,7 +131,7 @@ def handler(*args, **kwargs): # Always retry auth errors. if i == (max_retries - 1): # Exit the loop and wrap the authentication error. - raise _user_exceptions.FlyteAuthenticationException(_six.text_type(e)) + raise _user_exceptions.FlyteAuthenticationException(str(e)) cli_logger.error(f"Unauthenticated RPC error {e}, refreshing credentials and retrying\n") refresh_handler_fn = _get_refresh_handler(_creds_config.AUTH_MODE.get()) refresh_handler_fn(args[0]) diff --git a/flytekit/clis/flyte_cli/main.py b/flytekit/clis/flyte_cli/main.py index 438e91141f..82300a9c7f 100644 --- a/flytekit/clis/flyte_cli/main.py +++ b/flytekit/clis/flyte_cli/main.py @@ -20,23 +20,14 @@ from flytekit import __version__ from flytekit.clients import friendly as _friendly_client -from flytekit.clis.helpers import construct_literal_map_from_parameter_map as _construct_literal_map_from_parameter_map -from flytekit.clis.helpers import construct_literal_map_from_variable_map as _construct_literal_map_from_variable_map from flytekit.clis.helpers import hydrate_registration_parameters -from flytekit.clis.helpers import parse_args_into_dict as _parse_args_into_dict -from flytekit.common import launch_plan as _launch_plan_common -from flytekit.common import utils as _utils -from flytekit.common import workflow_execution as _workflow_execution_common -from flytekit.common.core import identifier as _identifier -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.tasks import task as _tasks_common -from flytekit.common.types import helpers as _type_helpers -from flytekit.common.utils import load_proto_from_file as _load_proto_from_file from flytekit.configuration import auth as _auth_config from flytekit.configuration import platform as _platform_config from flytekit.configuration import set_flyte_config_file -from flytekit.interfaces.data import data_proxy as _data_proxy -from flytekit.interfaces.data.data_proxy import Data +from flytekit.core import utils +from flytekit.core.context_manager import FlyteContextManager +from flytekit.exceptions import user as _user_exceptions +from flytekit.interfaces import cli_identifiers from flytekit.models import common as _common_models from flytekit.models import filters as _filters from flytekit.models import launch_plan as _launch_plan @@ -47,8 +38,6 @@ from flytekit.models.common import RawOutputDataConfig as _RawOutputDataConfig from flytekit.models.core import execution as _core_execution_models from flytekit.models.core import identifier as _core_identifier -from flytekit.models.execution import ExecutionMetadata as _ExecutionMetadata -from flytekit.models.execution import ExecutionSpec as _ExecutionSpec from flytekit.models.matchable_resource import ClusterResourceAttributes as _ClusterResourceAttributes from flytekit.models.matchable_resource import ExecutionClusterLabel as _ExecutionClusterLabel from flytekit.models.matchable_resource import ExecutionQueueAttributes as _ExecutionQueueAttributes @@ -118,20 +107,7 @@ def _get_io_string(literal_map, verbose=False): :param bool verbose: :rtype: Text """ - value_dict = _type_helpers.unpack_literal_map_to_sdk_object(literal_map) - if value_dict: - return "\n" + "\n".join( - "{:30}: {}".format( - k, - _prefix_lines( - "{:30} ".format(""), - v.verbose_string() if verbose else v.short_string(), - ), - ) - for k, v in value_dict.items() - ) - else: - return "(None)" + return str(literal_map) def _fetch_and_stringify_literal_map(path, verbose=False): @@ -140,12 +116,13 @@ def _fetch_and_stringify_literal_map(path, verbose=False): :param bool verbose: :rtype: Text """ - with _utils.AutoDeletingTempDir("flytecli") as tmp: + ctx = FlyteContextManager.current_context() + with utils.AutoDeletingTempDir("flytecli") as tmp: try: fname = tmp.get_named_tempfile("literalmap.pb") - _data_proxy.Data.get_data(path, fname) + ctx.file_access.get_data(path, fname) literal_map = _literals.LiteralMap.from_flyte_idl( - _utils.load_proto_from_file(_literals_pb2.LiteralMap, fname) + utils.load_proto_from_file(_literals_pb2.LiteralMap, fname) ) return _get_io_string(literal_map, verbose=verbose) except Exception: @@ -254,7 +231,7 @@ def _secho_one_execution(ex, urns_only): if not urns_only: _click.echo( "{:100} {:40} {:40}".format( - _tt(_identifier.WorkflowExecutionIdentifier.promote_from_model(ex.id)), + _tt(cli_identifiers.WorkflowExecutionIdentifier.promote_from_model(ex.id)), _tt(ex.id.name), _tt(ex.spec.launch_plan.name), ), @@ -263,7 +240,7 @@ def _secho_one_execution(ex, urns_only): _secho_workflow_status(ex.closure.phase) else: _click.echo( - "{:100}".format(_tt(_identifier.WorkflowExecutionIdentifier.promote_from_model(ex.id))), + "{:100}".format(_tt(cli_identifiers.WorkflowExecutionIdentifier.promote_from_model(ex.id))), nl=True, ) @@ -271,7 +248,7 @@ def _secho_one_execution(ex, urns_only): def _terminate_one_execution(client, urn, cause, shouldPrint=True): if shouldPrint: _click.echo("{:100} {:40}".format(_tt(urn), _tt(cause))) - client.terminate_execution(_identifier.WorkflowExecutionIdentifier.from_python_std(urn), cause) + client.terminate_execution(cli_identifiers.WorkflowExecutionIdentifier.from_python_std(urn), cause) def _update_one_launch_plan(client: _friendly_client.SynchronousFlyteClient, urn, state): @@ -279,7 +256,7 @@ def _update_one_launch_plan(client: _friendly_client.SynchronousFlyteClient, urn state = _launch_plan.LaunchPlanState.ACTIVE else: state = _launch_plan.LaunchPlanState.INACTIVE - client.update_launch_plan(_identifier.Identifier.from_python_std(urn), state) + client.update_launch_plan(cli_identifiers.Identifier.from_python_std(urn), state) _click.echo("Successfully updated {}".format(_tt(urn))) @@ -644,7 +621,7 @@ def parse_proto(filename, proto_class): idl_obj = split[-1] mod = _importlib.import_module(idl_module) idl = getattr(mod, idl_obj) - obj = _load_proto_from_file(idl, filename) + obj = utils.load_proto_from_file(idl, filename) jsonObj = MessageToJson(obj) @@ -734,7 +711,7 @@ def list_task_versions(project, domain, name, host, insecure, token, limit, show _click.echo( "{:50} {:40}".format( _tt(t.id.version), - _tt(_identifier.Identifier.promote_from_model(t.id)), + _tt(cli_identifiers.Identifier.promote_from_model(t.id)), ) ) @@ -760,57 +737,11 @@ def get_task(urn, host, insecure): _welcome_message() parent_ctx = _click.get_current_context(silent=True) client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) - t = client.get_task(_identifier.Identifier.from_python_std(urn)) + t = client.get_task(cli_identifiers.Identifier.from_python_std(urn)) _click.echo(_tt(t)) _click.echo("") -@_flyte_cli.command("launch-task", cls=_FlyteSubCommand) -@_project_option -@_domain_option -@_optional_name_option -@_assumable_iam_role_option -@_kubernetes_service_acct_option -@_host_option -@_insecure_option -@_urn_option -@_click.argument("task_args", nargs=-1, type=_click.UNPROCESSED) -def launch_task(project, domain, name, assumable_iam_role, kubernetes_service_account, host, insecure, urn, task_args): - """ - Kick off a single task execution. Note that the {project, domain, name} specified in the command line - will be for the execution. The project/domain for the task are specified in the urn. - - Use a -- to separate arguments to this cli, and arguments to the task. - e.g. - $ flyte-cli -h localhost:30081 -p flyteexamples -d development launch-task \ - -u tsk:flyteexamples:development:some-task:abc123 -- input=hi \ - other-input=123 moreinput=qwerty - - These arguments are then collected, and passed into the `task_args` variable as a Tuple[Text]. - Users should use the get-task command to ascertain the names of inputs to use. - """ - _welcome_message() - auth_role = _AuthRole(assumable_iam_role=assumable_iam_role, kubernetes_service_account=kubernetes_service_account) - - with _platform_config.URL.get_patcher(host), _platform_config.INSECURE.get_patcher(_tt(insecure)): - task_id = _identifier.Identifier.from_python_std(urn) - task = _tasks_common.SdkTask.fetch(task_id.project, task_id.domain, task_id.name, task_id.version) - - text_args = _parse_args_into_dict(task_args) - inputs = {} - for var_name, variable in task.interface.inputs.items(): - sdk_type = _type_helpers.get_sdk_type_from_literal_type(variable.type) - if var_name in text_args and text_args[var_name] is not None: - inputs[var_name] = sdk_type.from_string(text_args[var_name]).to_python_std() - - # TODO: Implement notification overrides - # TODO: Implement label overrides - # TODO: Implement annotation overrides - execution = task.launch(project, domain, inputs=inputs, name=name, auth_role=auth_role) - _click.secho("Launched execution: {}".format(_tt(execution.id)), fg="blue") - _click.echo("") - - ######################################################################################################################## # # Workflow Commands @@ -892,7 +823,7 @@ def list_workflow_versions(project, domain, name, host, insecure, token, limit, _click.echo( "{:50} {:40}".format( _tt(w.id.version), - _tt(_identifier.Identifier.promote_from_model(w.id)), + _tt(cli_identifiers.Identifier.promote_from_model(w.id)), ) ) @@ -918,7 +849,7 @@ def get_workflow(urn, host, insecure): _welcome_message() parent_ctx = _click.get_current_context(silent=True) client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) - _click.echo(client.get_workflow(_identifier.Identifier.from_python_std(urn))) + _click.echo(client.get_workflow(cli_identifiers.Identifier.from_python_std(urn))) # TODO: Print workflow pretty _click.echo("") @@ -1003,13 +934,13 @@ def list_active_launch_plans(project, domain, host, insecure, token, limit, show for lp in active_lps: if urns_only: - _click.echo("{:80}".format(_tt(_identifier.Identifier.promote_from_model(lp.id)))) + _click.echo("{:80}".format(_tt(cli_identifiers.Identifier.promote_from_model(lp.id)))) else: _click.echo( "{:30} {:50} {:80}".format( _render_schedule_expr(lp), _tt(lp.id.version), - _tt(_identifier.Identifier.promote_from_model(lp.id)), + _tt(cli_identifiers.Identifier.promote_from_model(lp.id)), ), ) @@ -1072,12 +1003,12 @@ def list_launch_plan_versions( ) for l in lp_list: if urns_only: - _click.echo(_tt(_identifier.Identifier.promote_from_model(l.id))) + _click.echo(_tt(cli_identifiers.Identifier.promote_from_model(l.id))) else: _click.echo( "{:50} {:80} ".format( _tt(l.id.version), - _tt(_identifier.Identifier.promote_from_model(l.id)), + _tt(cli_identifiers.Identifier.promote_from_model(l.id)), ), nl=False, ) @@ -1115,7 +1046,7 @@ def get_launch_plan(urn, host, insecure): _welcome_message() parent_ctx = _click.get_current_context(silent=True) client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) - _click.echo(_tt(client.get_launch_plan(_identifier.Identifier.from_python_std(urn)))) + _click.echo(_tt(client.get_launch_plan(cli_identifiers.Identifier.from_python_std(urn)))) # TODO: Print launch plan pretty _click.echo("") @@ -1167,50 +1098,6 @@ def update_launch_plan(state, host, insecure, urn=None): _update_one_launch_plan(client, urn=urn, state=state) -@_flyte_cli.command("execute-launch-plan", cls=_FlyteSubCommand) -@_project_option -@_domain_option -@_optional_name_option -@_host_option -@_insecure_option -@_urn_option -@_principal_option -@_verbose_option -@_watch_option -@_click.argument("lp_args", nargs=-1, type=_click.UNPROCESSED) -def execute_launch_plan(project, domain, name, host, insecure, urn, principal, verbose, watch, lp_args): - """ - Kick off a launch plan. Note that the {project, domain, name} specified in the command line - will be for the execution. The project/domain for the launch plan are specified in the urn. - - Use a -- to separate arguments to this cli, and arguments to the launch plan. - e.g. - $ flyte-cli -h localhost:30081 -p flyteexamples -d development execute-launch-plan \ - --verbose --principal=sdk-demo - -u lp:flyteexamples:development:some-workflow:abc123 -- input=hi \ - other-input=123 moreinput=qwerty - - These arguments are then collected, and passed into the `lp_args` variable as a Tuple[Text]. - Users should use the get-launch-plan command to ascertain the names of inputs to use. - """ - _welcome_message() - - with _platform_config.URL.get_patcher(host), _platform_config.INSECURE.get_patcher(_tt(insecure)): - lp_id = _identifier.Identifier.from_python_std(urn) - lp = _launch_plan_common.SdkLaunchPlan.fetch(lp_id.project, lp_id.domain, lp_id.name, lp_id.version) - - inputs = _construct_literal_map_from_parameter_map(lp.default_inputs, _parse_args_into_dict(lp_args)) - # TODO: Implement notification overrides - # TODO: Implement label overrides - # TODO: Implement annotation overrides - execution = lp.launch_with_literals(project, domain, inputs, name=name) - _click.secho("Launched execution: {}".format(_tt(execution.id)), fg="blue") - _click.echo("") - - if watch is True: - execution.wait_for_completion() - - ######################################################################################################################## # # Execution Commands @@ -1218,113 +1105,6 @@ def execute_launch_plan(project, domain, name, host, insecure, urn, principal, v ######################################################################################################################## -@_flyte_cli.command("watch-execution", cls=_FlyteSubCommand) -@_host_option -@_insecure_option -@_urn_option -def watch_execution(host, insecure, urn): - """ - Wait for an execution to complete. - - e.g. - $ flyte-cli -h localhost:30081 watch-execution -u ex:flyteexamples:development:abc123 - """ - _welcome_message() - parent_ctx = _click.get_current_context(silent=True) - client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) - ex_id = _identifier.WorkflowExecutionIdentifier.from_python_std(urn) - - execution = _workflow_execution_common.SdkWorkflowExecution.promote_from_model(client.get_execution(ex_id)) - - _click.echo("Waiting for the execution {} to complete ...".format(_tt(execution.id))) - - with _platform_config.URL.get_patcher(host), _platform_config.INSECURE.get_patcher(_tt(insecure)): - execution.wait_for_completion() - - -@_flyte_cli.command("relaunch-execution", cls=_FlyteSubCommand) -@_optional_project_option -@_optional_domain_option -@_optional_name_option -@_host_option -@_insecure_option -@_urn_option -@_optional_principal_option -@_verbose_option -@_click.argument("lp_args", nargs=-1, type=_click.UNPROCESSED) -def relaunch_execution(project, domain, name, host, insecure, urn, principal, verbose, lp_args): - """ - Relaunch a launch plan. - As with kicking off a launch plan (see execute-launch-plan), the project and domain will correspond to the new - execution to be run, and the project/domain used to find the existing execution will come from the URN. - This means you can re-run a development execution, in production, off of a staging launch-plan (in another project), - but beware that execution environment configurations can result in slower executions or permissions failures. - Therefore, it is recommended to re-run in the same environment as the original execution. By default, if the - project and domain are not specified, the existing project/domain will be used. - - When relaunching an execution, this will display the fixed inputs that it ran with (from the launch plan spec), - and handle the other inputs similar to how we handle initial launch plan execution, except that - all inputs now will have a default (the input of the execution being rerun). - - Use a -- to separate arguments to this cli, and arguments to the launch plan. - e.g. - $ flyte-cli -h localhost:30081 -p flyteexamples -d development execute-launch-plan \ - -u lp:flyteexamples:development:some-workflow:abc123 -- input=hi \ - other-input=123 moreinput=qwerty - - These arguments are then collected, and passed into the `lp_args` variable as a Tuple[Text]. - Users should use the get-execution and get-launch-plan commands to ascertain the names of inputs to use. - """ - _welcome_message() - parent_ctx = _click.get_current_context(silent=True) - client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) - - _click.echo("Relaunching execution {}\n".format(_tt(urn))) - existing_workflow_execution_identifier = _identifier.WorkflowExecutionIdentifier.from_python_std(urn) - e = client.get_execution(existing_workflow_execution_identifier) - - if project is None: - project = existing_workflow_execution_identifier.project - if domain is None: - domain = existing_workflow_execution_identifier.domain - if principal is None: - principal = e.spec.metadata.principal - - lp_model = client.get_launch_plan(e.spec.launch_plan) - expected_inputs = lp_model.closure.expected_inputs - - # Parse text inputs using the LP closure's parameter map to determine types. However, since all inputs are now - # optional (because we can default to the original execution's), we reduce first to bare Variables. - variable_map = {k: v.var for k, v in expected_inputs.parameters.items()} - parsed_text_args = _parse_args_into_dict(lp_args) - new_inputs = _construct_literal_map_from_variable_map(variable_map, parsed_text_args) - if len(new_inputs.literals) > 0: - _click.secho("\tNew Inputs: {}\n".format(_prefix_lines("\t\t", _get_io_string(new_inputs, verbose=verbose)))) - - # Construct new inputs from existing execution inputs and new inputs - inputs_dict = {} - for k in e.spec.inputs.literals.keys(): - if k in new_inputs.literals: - inputs_dict[k] = new_inputs.literals[k] - else: - inputs_dict[k] = e.spec.inputs.literals[k] - inputs = _literals.LiteralMap(literals=inputs_dict) - - if len(inputs_dict) > 0: - _click.secho( - "\tFinal Inputs for New Execution: {}\n".format( - _prefix_lines("\t\t", _get_io_string(inputs, verbose=verbose)) - ) - ) - - metadata = _ExecutionMetadata(mode=_ExecutionMetadata.ExecutionMode.MANUAL, principal=principal, nesting=0) - ex_spec = _ExecutionSpec(launch_plan=lp_model.id, inputs=inputs, metadata=metadata) - execution_identifier = client.create_execution(project=project, domain=domain, name=name, execution_spec=ex_spec) - execution_identifier = _identifier.WorkflowExecutionIdentifier.promote_from_model(execution_identifier) - _click.secho("Launched execution: {}".format(execution_identifier), fg="blue") - _click.echo("") - - @_flyte_cli.command("recover-execution", cls=_FlyteSubCommand) @_urn_option @_optional_name_option @@ -1356,10 +1136,10 @@ def recover_execution(urn, name, host, insecure): _click.echo("Recovering execution {}\n".format(_tt(urn))) - original_workflow_execution_identifier = _identifier.WorkflowExecutionIdentifier.from_python_std(urn) + original_workflow_execution_identifier = cli_identifiers.WorkflowExecutionIdentifier.from_python_std(urn) execution_identifier_resp = client.recover_execution(id=original_workflow_execution_identifier, name=name) - execution_identifier = _identifier.WorkflowExecutionIdentifier.promote_from_model(execution_identifier_resp) + execution_identifier = cli_identifiers.WorkflowExecutionIdentifier.promote_from_model(execution_identifier_resp) _click.secho("Launched execution: {}".format(execution_identifier), fg="blue") _click.echo("") @@ -1500,7 +1280,7 @@ def _render_workflow_execution(wf_execution, uri_to_message_map, show_io, verbos _click.echo( "\t{:15} {}".format( "Launch Plan:", - _tt(_identifier.Identifier.promote_from_model(wf_execution.spec.launch_plan)), + _tt(cli_identifiers.Identifier.promote_from_model(wf_execution.spec.launch_plan)), ) ) @@ -1675,7 +1455,7 @@ def _render_node_executions(client, node_execs, show_io, verbose, host, insecure "Subtasks:", "flyte-cli get-child-executions -h {host}{insecure} -u {urn}".format( host=host, - urn=_tt(_identifier.TaskExecutionIdentifier.promote_from_model(te.id)), + urn=_tt(cli_identifiers.TaskExecutionIdentifier.promote_from_model(te.id)), insecure=" --insecure" if insecure else "", ), ) @@ -1699,7 +1479,7 @@ def get_execution(urn, host, insecure, show_io, verbose): _welcome_message() parent_ctx = _click.get_current_context(silent=True) client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) - e = client.get_execution(_identifier.WorkflowExecutionIdentifier.from_python_std(urn)) + e = client.get_execution(cli_identifiers.WorkflowExecutionIdentifier.from_python_std(urn)) node_execs = _get_all_node_executions(client, workflow_execution_identifier=e.id) _render_node_executions(client, node_execs, show_io, verbose, host, insecure, wf_execution=e) @@ -1716,7 +1496,7 @@ def get_child_executions(urn, host, insecure, show_io, verbose): client = _friendly_client.SynchronousFlyteClient(host, insecure=insecure, root_cert_file=parent_ctx.obj["cacert"]) node_execs = _get_all_node_executions( client, - task_execution_identifier=_identifier.TaskExecutionIdentifier.from_python_std(urn), + task_execution_identifier=cli_identifiers.TaskExecutionIdentifier.from_python_std(urn), ) _render_node_executions(client, node_execs, show_io, verbose, host, insecure) @@ -1838,7 +1618,7 @@ def _extract_pair( f"Resource type found in proto file name [{resource_type}] invalid, " "must be 1 (task), 2 (workflow) or 3 (launch plan)" ) - entity = _load_proto_from_file(_resource_map[resource_type], object_file) + entity = utils.load_proto_from_file(_resource_map[resource_type], object_file) registerable_identifier, registerable_entity = hydrate_registration_parameters( resource_type, project, domain, version, entity ) @@ -2079,7 +1859,8 @@ def fast_register_files( version = version if version else digest full_remote_path = _get_additional_distribution_loc(additional_distribution_dir, version) - Data.put_data(compressed_source, full_remote_path) + ctx = FlyteContextManager.current_context() + ctx.file_access.put_data(compressed_source, full_remote_path) _click.secho(f"Uploaded compressed code archive {compressed_source} to {full_remote_path}", fg="green") def fast_register_task(entity: _GeneratedProtocolMessageType) -> _GeneratedProtocolMessageType: diff --git a/flytekit/clis/helpers.py b/flytekit/clis/helpers.py index c0ab7397c2..f87a7d2a11 100644 --- a/flytekit/clis/helpers.py +++ b/flytekit/clis/helpers.py @@ -7,34 +7,6 @@ from flyteidl.core import workflow_pb2 as _workflow_pb2 from flytekit.clis.sdk_in_container.serialize import _DOMAIN_PLACEHOLDER, _PROJECT_PLACEHOLDER, _VERSION_PLACEHOLDER -from flytekit.common.types.helpers import get_sdk_type_from_literal_type as _get_sdk_type_from_literal_type -from flytekit.models import literals as _literals - - -def construct_literal_map_from_variable_map(variable_dict, text_args): - """ - This function produces a map of Literals to use when creating an execution. It reads the required values from - a Variable map (presumably obtained from a launch plan), and then fills in the necessary inputs - from the click args. Click args will be strings, which will be parsed into their SDK types with each - SDK type's parse string method. - - :param dict[Text, flytekit.models.interface.Variable] variable_dict: - :param dict[Text, Text] text_args: - :rtype: flytekit.models.literals.LiteralMap - """ - inputs = {} - - for var_name, variable in variable_dict.items(): - # Check to see if it's passed from click - # If this is an input that has a default from the LP, it should've already been parsed into a string, - # and inserted into the default for this option, so it should still be here. - if var_name in text_args and text_args[var_name] is not None: - # the SDK type is also available from the sdk workflow object's saved user inputs but let's - # derive it here from the interface to be more rigorous. - sdk_type = _get_sdk_type_from_literal_type(variable.type) - inputs[var_name] = sdk_type.from_string(text_args[var_name]) - - return _literals.LiteralMap(literals=inputs) def parse_args_into_dict(input_arguments): @@ -49,36 +21,6 @@ def parse_args_into_dict(input_arguments): return {split_arg[0]: split_arg[1] for split_arg in [input_arg.split("=", 1) for input_arg in input_arguments]} -def construct_literal_map_from_parameter_map(parameter_map, text_args): - """ - Take a dictionary of Text to Text and construct a literal map using a ParameterMap as guidance. - Required input parameters must have an entry in the text arguments given. - Parameters with defaults will have those defaults filled in if missing from the text arguments. - - :param flytekit.models.interface.ParameterMap parameter_map: - :param dict[Text, Text] text_args: - :rtype: flytekit.models.literals.LiteralMap - """ - - # This function can be written by calling construct_literal_map_from_variable_map also, but not that much - # code is saved. - inputs = {} - for var_name, parameter in parameter_map.parameters.items(): - sdk_type = _get_sdk_type_from_literal_type(parameter.var.type) - if parameter.required: - if var_name in text_args and text_args[var_name] is not None: - inputs[var_name] = sdk_type.from_string(text_args[var_name]) - else: - raise Exception("Missing required parameter {}".format(var_name)) - else: - if var_name in text_args and text_args[var_name] is not None: - inputs[var_name] = sdk_type.from_string(text_args[var_name]) - else: - inputs[var_name] = parameter.default - - return _literals.LiteralMap(literals=inputs) - - def str2bool(str): """ bool('False') is True in Python, so we need to do some string parsing. Use the same words in ConfigParser diff --git a/flytekit/clis/sdk_in_container/basic_auth.py b/flytekit/clis/sdk_in_container/basic_auth.py index 8806c5c406..92612595ad 100644 --- a/flytekit/clis/sdk_in_container/basic_auth.py +++ b/flytekit/clis/sdk_in_container/basic_auth.py @@ -3,8 +3,8 @@ import requests as _requests -from flytekit.common.exceptions.user import FlyteAuthenticationException as _FlyteAuthenticationException from flytekit.configuration.creds import CLIENT_CREDENTIALS_SECRET as _CREDENTIALS_SECRET +from flytekit.exceptions.user import FlyteAuthenticationException _utf_8 = "utf-8" @@ -18,7 +18,7 @@ def get_secret(): secret = _CREDENTIALS_SECRET.get() if secret: return secret - raise _FlyteAuthenticationException("No secret could be found") + raise FlyteAuthenticationException("No secret could be found") def get_basic_authorization_header(client_id, client_secret): @@ -55,7 +55,7 @@ def get_token(token_endpoint, authorization_header, scope): response = _requests.post(token_endpoint, data=body, headers=headers) if response.status_code != 200: _logging.error("Non-200 ({}) received from IDP: {}".format(response.status_code, response.text)) - raise _FlyteAuthenticationException("Non-200 received from IDP") + raise FlyteAuthenticationException("Non-200 received from IDP") response = response.json() return response["access_token"], response["expires_in"] diff --git a/flytekit/clis/sdk_in_container/fast_register.py b/flytekit/clis/sdk_in_container/fast_register.py deleted file mode 100644 index 944ef502bf..0000000000 --- a/flytekit/clis/sdk_in_container/fast_register.py +++ /dev/null @@ -1,249 +0,0 @@ -import os as _os -from typing import List as _List - -import click - -from flytekit.clis.sdk_in_container.constants import CTX_DOMAIN, CTX_PACKAGES, CTX_PROJECT, CTX_TEST -from flytekit.common import utils as _utils -from flytekit.common.core import identifier as _identifier -from flytekit.common.tasks import sdk_runnable as _sdk_runnable_task -from flytekit.common.tasks import task as _task -from flytekit.configuration import sdk as _sdk_config -from flytekit.tools.fast_registration import compute_digest as _compute_digest -from flytekit.tools.fast_registration import get_additional_distribution_loc as _get_additional_distribution_loc -from flytekit.tools.fast_registration import upload_package as _upload_package -from flytekit.tools.module_loader import iterate_registerable_entities_in_order - - -def fast_register_all( - project: str, - domain: str, - pkgs: _List[str], - test: bool, - version: str, - source_dir: _os.PathLike, - dest_dir: _os.PathLike = None, -): - if test: - click.echo("Test switch enabled, not doing anything...") - - if not version: - digest = _compute_digest(source_dir) - else: - digest = version - remote_package_path = _upload_package(source_dir, digest, _sdk_config.FAST_REGISTRATION_DIR.get()) - - click.echo( - "Running task, workflow, and launch plan fast registration for {}, {}, {} with version {} and code dir {}".format( - project, domain, pkgs, digest, source_dir - ) - ) - - # m = module (i.e. python file) - # k = value of dir(m), type str - # o = object (e.g. SdkWorkflow) - for m, k, o in iterate_registerable_entities_in_order(pkgs): - name = _utils.fqdn(m.__name__, k, entity_type=o.resource_type) - o._id = _identifier.Identifier(o.resource_type, project, domain, name, digest) - - if test: - click.echo("Would fast register {:20} {}".format("{}:".format(o.entity_type_text), o.id.name)) - else: - click.echo("Fast registering {:20} {}".format("{}:".format(o.entity_type_text), o.id.name)) - _get_additional_distribution_loc(_sdk_config.FAST_REGISTRATION_DIR.get(), digest) - if isinstance(o, _sdk_runnable_task.SdkRunnableTask): - o.fast_register(project, domain, o.id.name, digest, remote_package_path, dest_dir) - else: - o.register(project, domain, o.id.name, digest) - - -def fast_register_tasks_only( - project: str, - domain: str, - pkgs: _List[str], - test: bool, - version: str, - source_dir: _os.PathLike, - dest_dir: _os.PathLike = None, -): - if test: - click.echo("Test switch enabled, not doing anything...") - - if not version: - digest = _compute_digest(source_dir) - else: - digest = version - remote_package_path = _upload_package(source_dir, digest, _sdk_config.FAST_REGISTRATION_DIR.get()) - - click.echo( - "Running task only fast registration for {}, {}, {} with version {} and code dir {}".format( - project, domain, pkgs, digest, source_dir - ) - ) - - # Discover all tasks by loading the module - for m, k, t in iterate_registerable_entities_in_order(pkgs, include_entities={_task.SdkTask}): - name = _utils.fqdn(m.__name__, k, entity_type=t.resource_type) - - if test: - click.echo("Would fast register task {:20} {}".format("{}:".format(t.entity_type_text), name)) - else: - click.echo("Fast registering task {:20} {}".format("{}:".format(t.entity_type_text), name)) - if isinstance(t, _sdk_runnable_task.SdkRunnableTask): - t.fast_register(project, domain, name, digest, remote_package_path, dest_dir) - else: - t.register(project, domain, name, digest) - - -@click.group("fast-register") -@click.option("--test", is_flag=True, help="Dry run, do not actually register with Admin") -@click.pass_context -def fast_register(ctx, test=None): - """ - Run fast registration steps for the Flyte entities in this container. This is an optimization to avoid the - conventional container build and upload cycle. This can be useful for fast iteration when making code changes. - If you do need to change the container itself (e.g. by adding a new dependency/import) you must rebuild and - upload a container. - - Caveats: Your flyte config must specify a fast registration dir like so: - [sdk] - fast_registration_dir=s3://my-s3-bucket/dir - - **and** ensure that the role specified in [auth] section of your config has read access to this remote location. - Furthermore, the role you assume to call fast-register must have **write** permission to this remote location. - - Run with the --test switch for a dry run to see what will be registered. A default launch plan will also be - created, if a role can be found in the environment variables. - """ - - ctx.obj[CTX_TEST] = test - - -@click.command("tasks") -@click.option( - "--source-dir", - type=str, - help="The root dir of the code that should be uploaded for fast registration.", - required=True, -) -@click.option( - "--dest-dir", - type=str, - help="[Optional] The output directory of code which is downloaded during fast registration. " - "If the current working directory at the time of installation is not desired", -) -@click.option( - "-v", - "--version", - type=str, - help="Version to register tasks with. This is normally computed deterministically from your code, " - "but you can override here.", -) -@click.pass_context -def tasks(ctx, source_dir, dest_dir=None, version=None): - """ - Only fast register tasks. - - For example, consider a sample directory where tasks defined in workflows/ imports code from util/ like so: - - \b - $ tree /root/code/ - /root/code/ - ├── Dockerfile - ├── Makefile - ├── README.md - ├── conf.py - ├── notebook.config - ├── workflows - │   ├── __init__.py - │   ├── compose - │   │   ├── README.md - │   │   ├── __init__.py - │   │   ├── a_workflow.py - │   │   ├── b_workflow.py - ├── util - │   ├── __init__.py - │   ├── shared_task_code.py - ├── requirements.txt - ├── flyte.config - - Your source dir will need to be /root/code/ rather than the workflow packages dir /root/code/workflows you might - have specified in your flyte.config because all of the code your workflows depends on needs to be encapsulated in - `source_dir`, like so: - - pyflyte -p myproject -d development fast-register tasks --source-dir /root/code/ - - """ - project = ctx.obj[CTX_PROJECT] - domain = ctx.obj[CTX_DOMAIN] - test = ctx.obj[CTX_TEST] - pkgs = ctx.obj[CTX_PACKAGES] - - fast_register_tasks_only(project, domain, pkgs, test, version, source_dir, dest_dir) - - -@click.command("workflows") -@click.option( - "--source-dir", - type=str, - help="The root dir of the code that should be uploaded for fast registration.", - required=True, -) -@click.option( - "--dest-dir", - type=str, - help="[Optional] The output directory of code which is downloaded during fast registration. " - "If the current working directory at the time of installation is not desired", -) -@click.option( - "-v", - "--version", - type=str, - help="Version to register entities with. This is normally computed deterministically from your code, " - "but you can override here.", -) -@click.pass_context -def workflows(ctx, source_dir, dest_dir=None, version=None): - """ - Fast register both tasks and workflows. Also create and register a default launch plan for all workflows. - The `source_dir` param should point to the root directory of your project that contains all of your working code. - - For example, consider a sample directory structure where code in workflows/ imports code from util/ like so: - - \b - $ tree /root/code/ - /root/code/ - ├── Dockerfile - ├── Makefile - ├── README.md - ├── conf.py - ├── notebook.config - ├── workflows - │   ├── __init__.py - │   ├── compose - │   │   ├── README.md - │   │   ├── __init__.py - │   │   ├── a_workflow.py - │   │   ├── b_workflow.py - ├── util - │   ├── __init__.py - │   ├── shared_workflow_code.py - ├── requirements.txt - ├── flyte.config - - Your source dir will need to be /root/code/ rather than the workflow packages dir /root/code/workflows you might - have specified in your flyte.config because all of the code your workflows depends on needs to be encapsulated in - `source_dir`, like so: - - pyflyte -p myproject -d development fast-register workflows --source-dir /root/code/ - """ - project = ctx.obj[CTX_PROJECT] - domain = ctx.obj[CTX_DOMAIN] - test = ctx.obj[CTX_TEST] - pkgs = ctx.obj[CTX_PACKAGES] - - fast_register_all(project, domain, pkgs, test, version, source_dir, dest_dir) - - -fast_register.add_command(tasks) -fast_register.add_command(workflows) diff --git a/flytekit/clis/sdk_in_container/launch_plan.py b/flytekit/clis/sdk_in_container/launch_plan.py deleted file mode 100644 index e367fa08d8..0000000000 --- a/flytekit/clis/sdk_in_container/launch_plan.py +++ /dev/null @@ -1,276 +0,0 @@ -import logging as _logging -import os as _os - -import click -import six as _six - -from flytekit.clis.helpers import construct_literal_map_from_parameter_map as _construct_literal_map_from_parameter_map -from flytekit.clis.sdk_in_container import constants as _constants -from flytekit.clis.sdk_in_container.constants import ( - CTX_DOMAIN, - CTX_PROJECT, - CTX_VERSION, - domain_option, - project_option, - version_option, -) -from flytekit.common import utils as _utils -from flytekit.common.launch_plan import SdkLaunchPlan as _SdkLaunchPlan -from flytekit.configuration.internal import DOMAIN as _DOMAIN -from flytekit.configuration.internal import IMAGE as _IMAGE -from flytekit.configuration.internal import PROJECT as _PROJECT -from flytekit.configuration.internal import VERSION as _VERSION -from flytekit.configuration.internal import look_up_version_from_image_tag as _look_up_version_from_image_tag -from flytekit.models import launch_plan as _launch_plan_model -from flytekit.models.core import identifier as _identifier -from flytekit.tools.module_loader import iterate_registerable_entities_in_order - - -class LaunchPlanAbstractGroup(click.Group): - """ - This class iterates over the workflow folders and loads all workflows that are implemented via the programming - model. - """ - - def __init__(self, name, **attrs): - super(LaunchPlanAbstractGroup, self).__init__(name, commands=None, **attrs) - - def list_commands(self, ctx): - commands = [] - lps = {} - pkgs = ctx.obj[_constants.CTX_PACKAGES] - # Discover all launch plans by loading the modules - for m, k, lp in iterate_registerable_entities_in_order( - pkgs, include_entities={_SdkLaunchPlan}, detect_unreferenced_entities=False - ): - safe_name = _utils.fqdn(m.__name__, k, entity_type=lp.resource_type) - commands.append(safe_name) - lps[safe_name] = lp - - ctx.obj["lps"] = lps - commands.sort() - - return commands - - def get_command(self, ctx, lp_argument): - # Get the launch plan object in one of two ways. If get_command is being called by the list function - # then it should have been cached in the context. - # If we are actually running the command, then it won't have been cached and we'll have to load everything again - launch_plan = None - pkgs = ctx.obj[_constants.CTX_PACKAGES] - - if "lps" in ctx.obj: - launch_plan = ctx.obj["lps"][lp_argument] - else: - for m, k, lp in iterate_registerable_entities_in_order( - pkgs, - include_entities={_SdkLaunchPlan}, - detect_unreferenced_entities=False, - ): - safe_name = _utils.fqdn(m.__name__, k, entity_type=lp.resource_type) - if lp_argument == safe_name: - launch_plan = lp - - if launch_plan is None: - raise Exception("Could not load launch plan {}".format(lp_argument)) - - launch_plan._id = _identifier.Identifier( - _identifier.ResourceType.LAUNCH_PLAN, - ctx.obj[_constants.CTX_PROJECT], - ctx.obj[_constants.CTX_DOMAIN], - lp_argument, - ctx.obj[_constants.CTX_VERSION], - ) - return self._get_command(ctx, launch_plan, lp_argument) - - def _get_command(self, ctx, lp, cmd_name): - """ - :param ctx: - :param flytekit.common.launch_plan.SdkLaunchPlan lp: - :rtype: click.Command - """ - pass - - -class LaunchPlanExecuteGroup(LaunchPlanAbstractGroup): - def _get_command(self, ctx, lp, cmd_name): - """ - This function returns the function that click will actually use to execute a specific launch plan. It also - stores the launch plan python object and the command name in the closure. - - :param ctx: - :param flytekit.common.launch_plan.SdkLaunchPlan lp: - :param Text cmd_name: The name of the launch plan, as passed in from the abstract class - """ - - def _execute_lp(**kwargs): - for input_name in _six.iterkeys(kwargs): - if isinstance(kwargs[input_name], tuple): - kwargs[input_name] = list(kwargs[input_name]) - - inputs = _construct_literal_map_from_parameter_map(lp.default_inputs, kwargs) - execution = lp.execute_with_literals( - ctx.obj[_constants.CTX_PROJECT], - ctx.obj[_constants.CTX_DOMAIN], - literal_inputs=inputs, - notification_overrides=ctx.obj.get(_constants.CTX_NOTIFICATIONS, None), - ) - click.echo( - click.style( - "Workflow scheduled, execution_id={}".format(_six.text_type(execution.id)), - fg="blue", - ) - ) - - command = click.Command(name=cmd_name, callback=_execute_lp) - - # Iterate through the workflow's inputs - for var_name in sorted(lp.default_inputs.parameters): - param = lp.default_inputs.parameters[var_name] - # TODO: Figure out how to better handle the fact that we want strings to parse, - # but we probably shouldn't have click say that that's the type on the CLI. - help_msg = "{} Type: {}".format( - _six.text_type(param.var.description), _six.text_type(param.var.type) - ).strip() - - if param.required: - # If it's a required input, add the required flag - wrapper = click.option( - "--{}".format(var_name), - required=True, - type=_six.text_type, - help=help_msg, - ) - else: - # If it's not a required input, it should have a default - # Use to_python_std so that the text of the default ends up being parseable, if not, the click - # arg would look something like 'Integer(10)'. If the user specified '11' on the cli, then - # we'd get '11' and then we'd need annoying logic to differentiate between the default text - # and user text. - default = param.default.to_python_std() - wrapper = click.option( - "--{}".format(var_name), - default="{}".format(_six.text_type(default)), - type=_six.text_type, - help="{}. Default: {}".format(help_msg, _six.text_type(default)), - ) - - command = wrapper(command) - - return command - - -@click.group("lp") -@project_option -@domain_option -@version_option -@click.pass_context -def launch_plans(ctx, project, domain, version): - """ - Launch plan control group, including executions - """ - - version = version or _look_up_version_from_image_tag(_IMAGE.get()) - if not version: - raise click.UsageError("Could not find image from config, please specify a value for ``--version``") - - ctx.obj[CTX_PROJECT] = project - ctx.obj[CTX_DOMAIN] = domain - ctx.obj[CTX_VERSION] = version - _os.environ[_PROJECT.env_var] = project - _os.environ[_DOMAIN.env_var] = domain - _os.environ[_VERSION.env_var] = version - - -@click.group("execute", cls=LaunchPlanExecuteGroup) -@click.pass_context -def execute_launch_plan(ctx): - """ - Execute launch plans found in this container - """ - pass - - -def activate_all_impl(project, domain, version, pkgs, ignore_schedules=False): - # TODO: This should be a transaction to ensure all or none are updated - # TODO: We should optionally allow deactivation of missing launch plans - - # Discover all launch plans by loading the modules - _logging.info(f"Setting this version's {version} launch plans active in {project} {domain}") - for m, k, lp in iterate_registerable_entities_in_order( - pkgs, include_entities={_SdkLaunchPlan}, detect_unreferenced_entities=False - ): - lp._id = _identifier.Identifier( - _identifier.ResourceType.LAUNCH_PLAN, - project, - domain, - _utils.fqdn(m.__name__, k, entity_type=lp.resource_type), - version, - ) - if not (lp.is_scheduled and ignore_schedules): - _logging.info(f"Setting active {_utils.fqdn(m.__name__, k, entity_type=lp.resource_type)}") - lp.update(_launch_plan_model.LaunchPlanState.ACTIVE) - - -@click.command("activate-all-schedules") -@click.option( - "-v", - "--version", - type=str, - help="Version to register tasks with. This is normally parsed from the" "image, but you can override here.", -) -@click.pass_context -def activate_all_schedules(ctx, version=None): - """ - THIS COMMAND IS DEPRECATED. PLEASE USE activate-all - - The behavior of this command is identical to activate-all. - """ - click.secho( - "activate-all-schedules is deprecated, please use activate-all instead.", - color="yellow", - ) - project = ctx.obj[_constants.CTX_PROJECT] - domain = ctx.obj[_constants.CTX_DOMAIN] - pkgs = ctx.obj[_constants.CTX_PACKAGES] - version = version or ctx.obj[_constants.CTX_VERSION] or _look_up_version_from_image_tag(_IMAGE.get()) - activate_all_impl(project, domain, version, pkgs) - - -@click.command("activate-all") -@click.option( - "-v", - "--version", - type=str, - help="Version to register tasks with. This is normally parsed from the" "image, but you can override here.", -) -@click.option( - "--ignore-schedules", - is_flag=True, - help="Activate all except for launch plans with schedules.", -) -@click.pass_context -def activate_all(ctx, version=None, ignore_schedules=False): - """ - This command will activate all found launch plans at the given version. If there are existing - active launch plans that collide on project, domain, and name, but differ on version, those will be - deactivated in favor of the version specified in this command. If a launch plan is associated with a schedule, - the schedule will also be deactivated or activated as appropriate. - - Note: - 1. Currently, this is not a transaction. Therefore, if the command fails, it is possible that some schedules - have been updated. - 2. If a launch plan is scheduled on an older version for a given project, domain, and name AND there is not a - matching scheduled launch plan found when running this command, the existing schedule will remain active - until it is manually disabled. - """ - project = ctx.obj[_constants.CTX_PROJECT] - domain = ctx.obj[_constants.CTX_DOMAIN] - pkgs = ctx.obj[_constants.CTX_PACKAGES] - version = version or ctx.obj[_constants.CTX_VERSION] or _look_up_version_from_image_tag(_IMAGE.get()) - activate_all_impl(project, domain, version, pkgs, ignore_schedules=ignore_schedules) - - -launch_plans.add_command(execute_launch_plan) -launch_plans.add_command(activate_all_schedules) -launch_plans.add_command(activate_all) diff --git a/flytekit/clis/sdk_in_container/pyflyte.py b/flytekit/clis/sdk_in_container/pyflyte.py index b23fa5d121..e0d3859122 100644 --- a/flytekit/clis/sdk_in_container/pyflyte.py +++ b/flytekit/clis/sdk_in_container/pyflyte.py @@ -5,12 +5,9 @@ import click from flytekit.clis.sdk_in_container.constants import CTX_PACKAGES -from flytekit.clis.sdk_in_container.fast_register import fast_register from flytekit.clis.sdk_in_container.init import init -from flytekit.clis.sdk_in_container.launch_plan import launch_plans from flytekit.clis.sdk_in_container.local_cache import local_cache from flytekit.clis.sdk_in_container.package import package -from flytekit.clis.sdk_in_container.register import register from flytekit.clis.sdk_in_container.serialize import serialize from flytekit.configuration import internal as _internal_config from flytekit.configuration import platform as _platform_config @@ -107,10 +104,7 @@ def update_configuration_file(config_file_path): click.secho("Flyte Admin URL {}".format(_URL.get()), fg="green") -main.add_command(register) -main.add_command(fast_register) main.add_command(serialize) -main.add_command(launch_plans) main.add_command(package) main.add_command(local_cache) main.add_command(init) diff --git a/flytekit/clis/sdk_in_container/register.py b/flytekit/clis/sdk_in_container/register.py deleted file mode 100644 index 876629b851..0000000000 --- a/flytekit/clis/sdk_in_container/register.py +++ /dev/null @@ -1,146 +0,0 @@ -import logging as _logging -import os as _os - -import click - -from flytekit.clis.sdk_in_container.constants import ( - CTX_DOMAIN, - CTX_PACKAGES, - CTX_PROJECT, - CTX_TEST, - CTX_VERSION, - domain_option, - project_option, - version_option, -) -from flytekit.common import utils as _utils -from flytekit.common.core import identifier as _identifier -from flytekit.common.tasks import task as _task -from flytekit.configuration.internal import DOMAIN as _DOMAIN -from flytekit.configuration.internal import IMAGE as _IMAGE -from flytekit.configuration.internal import PROJECT as _PROJECT -from flytekit.configuration.internal import VERSION as _VERSION -from flytekit.configuration.internal import look_up_version_from_image_tag as _look_up_version_from_image_tag -from flytekit.tools.module_loader import iterate_registerable_entities_in_order - - -def register_all(project, domain, pkgs, test, version): - if test: - click.echo("Test switch enabled, not doing anything...") - click.echo( - "Running task, workflow, and launch plan registration for {}, {}, {} with version {}".format( - project, domain, pkgs, version - ) - ) - - # m = module (i.e. python file) - # k = value of dir(m), type str - # o = object (e.g. SdkWorkflow) - for m, k, o in iterate_registerable_entities_in_order(pkgs): - name = _utils.fqdn(m.__name__, k, entity_type=o.resource_type) - _logging.debug("Found module {}\n K: {} Instantiated in {}".format(m, k, o._instantiated_in)) - o._id = _identifier.Identifier(o.resource_type, project, domain, name, version) - - if test: - click.echo("Would register {:20} {}".format("{}:".format(o.entity_type_text), o.id.name)) - else: - click.echo("Registering {:20} {}".format("{}:".format(o.entity_type_text), o.id.name)) - o.register(project, domain, o.id.name, version) - - -def register_tasks_only(project, domain, pkgs, test, version): - if test: - click.echo("Test switch enabled, not doing anything...") - - click.echo("Running task only registration for {}, {}, {} with version {}".format(project, domain, pkgs, version)) - - # Discover all tasks by loading the module - for m, k, t in iterate_registerable_entities_in_order(pkgs, include_entities={_task.SdkTask}): - name = _utils.fqdn(m.__name__, k, entity_type=t.resource_type) - - if test: - click.echo("Would register task {:20} {}".format("{}:".format(t.entity_type_text), name)) - else: - click.echo("Registering task {:20} {}".format("{}:".format(t.entity_type_text), name)) - t.register(project, domain, name, version) - - -@click.group("register") -@project_option -@domain_option -@version_option -# --pkgs on the register group is DEPRECATED, use same arg on pyflyte.main instead -@click.option( - "--pkgs", - multiple=True, - help="DEPRECATED. This arg can only be used before the 'register' keyword", -) -@click.option("--test", is_flag=True, help="Dry run, do not actually register with Admin") -@click.pass_context -def register(ctx, project, domain, version, pkgs=None, test=None): - """ - Run registration steps for the workflows in this container. - - Run with the --test switch for a dry run to see what will be registered. A default launch plan will also be - created, if a role can be found in the environment variables. - """ - if pkgs: - raise click.UsageError("--pkgs must now be specified before the 'register' keyword on the command line") - - version = version or _look_up_version_from_image_tag(_IMAGE.get()) - if not version: - raise click.UsageError("Could not find image from config, please specify a value for ``--version``") - - ctx.obj[CTX_PROJECT] = project - ctx.obj[CTX_DOMAIN] = domain - ctx.obj[CTX_VERSION] = version - ctx.obj[CTX_TEST] = test - _os.environ[_PROJECT.env_var] = project - _os.environ[_DOMAIN.env_var] = domain - _os.environ[_VERSION.env_var] = ctx.obj[CTX_VERSION] - - -@click.command("tasks") -@click.option( - "-v", - "--version", - type=str, - help="Version to register tasks with. This is normally parsed from the" "image, but you can override here.", -) -@click.pass_context -def tasks(ctx, version=None): - """ - Only register tasks. - """ - project = ctx.obj[CTX_PROJECT] - domain = ctx.obj[CTX_DOMAIN] - test = ctx.obj[CTX_TEST] - pkgs = ctx.obj[CTX_PACKAGES] - - version = version or ctx.obj[CTX_VERSION] or _look_up_version_from_image_tag(_IMAGE.get()) - register_tasks_only(project, domain, pkgs, test, version) - - -@click.command("workflows") -@click.option( - "-v", - "--version", - type=str, - help="Version to register tasks with. This is normally parsed from the" "image, but you can override here.", -) -@click.pass_context -def workflows(ctx, version=None): - """ - Register both tasks and workflows. Also create and register a default launch plan for all workflows. - """ - project = ctx.obj[CTX_PROJECT] - domain = ctx.obj[CTX_DOMAIN] - test = ctx.obj[CTX_TEST] - pkgs = ctx.obj[CTX_PACKAGES] - - version = version or ctx.obj[CTX_VERSION] or _look_up_version_from_image_tag(_IMAGE.get()) - register_all(project, domain, pkgs, test, version) - - -register.add_command(tasks) -register.add_command(workflows) diff --git a/flytekit/clis/sdk_in_container/serialize.py b/flytekit/clis/sdk_in_container/serialize.py index f02526a5dd..5fd68aebc4 100644 --- a/flytekit/clis/sdk_in_container/serialize.py +++ b/flytekit/clis/sdk_in_container/serialize.py @@ -14,24 +14,21 @@ import flytekit as _flytekit from flytekit.clis.sdk_in_container.constants import CTX_PACKAGES -from flytekit.common import utils as _utils -from flytekit.common.core import identifier as _identifier -from flytekit.common.exceptions.scopes import system_entry_point -from flytekit.common.exceptions.user import FlyteValidationException -from flytekit.common.tasks import task as _sdk_task -from flytekit.common.translator import get_serializable -from flytekit.common.utils import write_proto_to_file as _write_proto_to_file from flytekit.configuration import internal as _internal_config from flytekit.core import context_manager as flyte_context from flytekit.core.base_task import PythonTask from flytekit.core.launch_plan import LaunchPlan from flytekit.core.workflow import WorkflowBase +from flytekit.exceptions.scopes import system_entry_point +from flytekit.exceptions.user import FlyteValidationException from flytekit.models import launch_plan as _launch_plan_models from flytekit.models import task as task_models from flytekit.models.admin import workflow as admin_workflow_models +from flytekit.models.core import identifier as _identifier from flytekit.tools.fast_registration import compute_digest as _compute_digest from flytekit.tools.fast_registration import filter_tar_file_fn as _filter_tar_file_fn -from flytekit.tools.module_loader import iterate_registerable_entities_in_order +from flytekit.tools.module_loader import trigger_loading +from flytekit.tools.translator import get_serializable # Identifier fields use placeholders for registration-time substitution. # Additional fields, such as auth and the raw output data prefix have more complex structures @@ -59,42 +56,6 @@ class SerializationMode(_Enum): FAST = 1 -@system_entry_point -def serialize_tasks_only(pkgs, folder=None): - """ - :param list[Text] pkgs: - :param Text folder: - - :return: - """ - # m = module (i.e. python file) - # k = value of dir(m), type str - # o = object (e.g. SdkWorkflow) - loaded_entities = [] - for m, k, o in iterate_registerable_entities_in_order(pkgs, include_entities={_sdk_task.SdkTask}): - name = _utils.fqdn(m.__name__, k, entity_type=o.resource_type) - _logging.debug("Found module {}\n K: {} Instantiated in {}".format(m, k, o._instantiated_in)) - o._id = _identifier.Identifier( - o.resource_type, _PROJECT_PLACEHOLDER, _DOMAIN_PLACEHOLDER, name, _VERSION_PLACEHOLDER - ) - loaded_entities.append(o) - - zero_padded_length = _determine_text_chars(len(loaded_entities)) - for i, entity in enumerate(loaded_entities): - serialized = entity.serialize() - fname_index = str(i).zfill(zero_padded_length) - fname = "{}_{}.pb".format(fname_index, entity._id.name) - click.echo(" Writing {} to\n {}".format(entity._id, fname)) - if folder: - fname = _os.path.join(folder, fname) - _write_proto_to_file(serialized, fname) - - identifier_fname = "{}_{}.identifier.pb".format(fname_index, entity._id.name) - if folder: - identifier_fname = _os.path.join(folder, identifier_fname) - _write_proto_to_file(entity._id.to_flyte_idl(), identifier_fname) - - def _should_register_with_admin(entity) -> bool: """ This is used in the code below. The translator.py module produces lots of objects (namely nodes and BranchNodes) @@ -140,8 +101,9 @@ def get_registrable_entities(ctx: flyte_context.FlyteContext) -> typing.List: serializable_tasks: typing.List[task_models.TaskSpec] = [ entity for entity in entities_to_be_serialized if isinstance(entity, task_models.TaskSpec) ] - # Detect if any of the tasks is duplicated. Duplicate tasks are defined as having the same metadata identifiers - # (see :py:class:`flytekit.common.core.identifier.Identifier`). Duplicate tasks are considered invalid at registration + # Detect if any of the tasks is duplicated. Duplicate tasks are defined as having the same + # metadata identifiers (see :py:class:`flytekit.common.core.identifier.Identifier`). Duplicate + # tasks are considered invalid at registration # time and usually indicate user error, so we catch this common mistake at serialization time. duplicate_tasks = _find_duplicate_tasks(serializable_tasks) if len(duplicate_tasks) > 0: @@ -211,9 +173,6 @@ def serialize_all( :param flytekit_virtualenv_root: The full path of the virtual env in the container. """ - # m = module (i.e. python file) - # k = value of dir(m), type str - # o = object (e.g. SdkWorkflow) env = { _internal_config.CONFIGURATION_PATH.env_var: config_path if config_path @@ -242,29 +201,9 @@ def serialize_all( ) ctx = flyte_context.FlyteContextManager.current_context().with_serialization_settings(serialization_settings) with flyte_context.FlyteContextManager.with_context(ctx) as ctx: - old_style_entities = [] - # This first for loop is for legacy API entities - SdkTask, SdkWorkflow, etc. The _get_entity_to_module - # function that this iterate calls only works on legacy objects - for m, k, o in iterate_registerable_entities_in_order(pkgs, local_source_root=local_source_root): - name = _utils.fqdn(m.__name__, k, entity_type=o.resource_type) - _logging.debug("Found module {}\n K: {} Instantiated in {}".format(m, k, o._instantiated_in)) - o._id = _identifier.Identifier( - o.resource_type, _PROJECT_PLACEHOLDER, _DOMAIN_PLACEHOLDER, name, _VERSION_PLACEHOLDER - ) - old_style_entities.append(o) - - serialized_old_style_entities = [] - for entity in old_style_entities: - if entity.has_registered: - _logging.info(f"Skipping entity {entity.id} because already registered") - continue - serialized_old_style_entities.append(entity.serialize()) - + trigger_loading(pkgs, local_source_root=local_source_root) click.echo(f"Found {len(flyte_context.FlyteEntities.entities)} tasks/workflows") - - new_api_model_values = get_registrable_entities(ctx) - - loaded_entities = serialized_old_style_entities + new_api_model_values + loaded_entities = get_registrable_entities(ctx) if folder is None: folder = "." persist_registrable_entities(loaded_entities, folder) @@ -349,18 +288,6 @@ def serialize(ctx, image, local_source_root, in_container_config_path, in_contai ctx.obj[CTX_PYTHON_INTERPRETER] = sys.executable -@click.command("tasks") -@click.option("-f", "--folder", type=click.Path(exists=True)) -@click.pass_context -def tasks(ctx, folder=None): - pkgs = ctx.obj[CTX_PACKAGES] - - if folder: - click.echo(f"Writing output to {folder}") - - serialize_tasks_only(pkgs, folder) - - @click.command("workflows") # For now let's just assume that the directory needs to exist. If you're docker run -v'ing, docker will create the # directory for you so it shouldn't be a problem. @@ -423,7 +350,5 @@ def fast_workflows(ctx, folder=None): fast.add_command(fast_workflows) - -serialize.add_command(tasks) serialize.add_command(workflows) serialize.add_command(fast) diff --git a/flytekit/common/component_nodes.py b/flytekit/common/component_nodes.py deleted file mode 100644 index ea39a28dea..0000000000 --- a/flytekit/common/component_nodes.py +++ /dev/null @@ -1,157 +0,0 @@ -import logging as _logging - -from flytekit.common import sdk_bases as _sdk_bases -from flytekit.common.exceptions import system as _system_exceptions -from flytekit.models.core import workflow as _workflow_model - - -class SdkTaskNode(_workflow_model.TaskNode, metaclass=_sdk_bases.ExtendedSdkType): - def __init__(self, sdk_task): - """ - :param flytekit.common.tasks.task.SdkTask sdk_task: - """ - self._sdk_task = sdk_task - super(SdkTaskNode, self).__init__(None) - - @property - def reference_id(self): - """ - A globally unique identifier for the task. - :rtype: flytekit.models.core.identifier.Identifier - """ - return self._sdk_task.id - - @property - def sdk_task(self): - """ - :rtype: flytekit.common.tasks.task.SdkTask - """ - return self._sdk_task - - @classmethod - def promote_from_model(cls, base_model, tasks): - """ - Takes the idl wrapper for a TaskNode and returns the hydrated Flytekit object for it by fetching it from the - engine. - - :param flytekit.models.core.workflow.TaskNode base_model: - :param dict[flytekit.models.core.identifier.Identifier, flytekit.models.task.TaskTemplate] tasks: - :rtype: SdkTaskNode - """ - from flytekit.common.tasks import task as _task - - if base_model.reference_id in tasks: - t = tasks[base_model.reference_id] - _logging.info(f"Found existing task template for {t.id}, will not retrieve from Admin") - sdk_task = _task.SdkTask.promote_from_model(t) - sdk_task._has_registered = True - return cls(sdk_task) - - # If not found, fetch it from Admin - _logging.debug("Fetching task template for {} from Admin".format(base_model.reference_id)) - project = base_model.reference_id.project - domain = base_model.reference_id.domain - name = base_model.reference_id.name - version = base_model.reference_id.version - sdk_task = _task.SdkTask.fetch(project, domain, name, version) - return cls(sdk_task) - - -class SdkWorkflowNode(_workflow_model.WorkflowNode, metaclass=_sdk_bases.ExtendedSdkType): - def __init__(self, sdk_workflow=None, sdk_launch_plan=None): - """ - :param flytekit.common.workflow.SdkWorkflow sdk_workflow: - :param flytekit.common.launch_plan.SdkLaunchPlan sdk_launch_plan: - """ - if sdk_workflow and sdk_launch_plan: - raise _system_exceptions.FlyteSystemException( - "SdkWorkflowNode cannot be called with both a workflow and " - "a launchplan specified, please pick one. WF: {} LP: {}", - sdk_workflow, - sdk_launch_plan, - ) - - self._sdk_workflow = sdk_workflow - self._sdk_launch_plan = sdk_launch_plan - sdk_wf_id = sdk_workflow.id if sdk_workflow else None - sdk_lp_id = sdk_launch_plan.id if sdk_launch_plan else None - super(SdkWorkflowNode, self).__init__(launchplan_ref=sdk_lp_id, sub_workflow_ref=sdk_wf_id) - - def __repr__(self): - """ - :rtype: Text - """ - if self.sdk_workflow is not None: - return "SdkWorkflowNode with workflow: {}".format(self.sdk_workflow) - return "SdkWorkflowNode with launch plan: {}".format(self.sdk_launch_plan) - - @property - def launchplan_ref(self): - """ - [Optional] A globally unique identifier for the launch plan. Should map to Admin. - :rtype: flytekit.models.core.identifier.Identifier - """ - return self._sdk_launch_plan.id if self._sdk_launch_plan else None - - @property - def sub_workflow_ref(self): - """ - [Optional] Reference to a subworkflow, that should be defined with the compiler context. - :rtype: flytekit.models.core.identifier.Identifier - """ - return self._sdk_workflow.id if self._sdk_workflow else None - - @property - def sdk_launch_plan(self): - """ - :rtype: flytekit.common.launch_plan.SdkLaunchPlan - """ - return self._sdk_launch_plan - - @property - def sdk_workflow(self): - """ - :rtype: flytekit.common.workflow.SdkWorkflow - """ - return self._sdk_workflow - - @classmethod - def promote_from_model(cls, base_model, sub_workflows, tasks): - """ - :param flytekit.models.core.workflow.WorkflowNode base_model: - :param dict[flytekit.models.core.identifier.Identifier, flytekit.models.core.workflow.WorkflowTemplate] - sub_workflows: - :param dict[flytekit.models.core.identifier.Identifier, flytekit.models.task.TaskTemplate] tasks: - :rtype: SdkWorkflowNode - """ - # put the import statement here to prevent circular dependency error - from flytekit.common import launch_plan as _launch_plan - from flytekit.common import workflow as _workflow - - project = base_model.reference.project - domain = base_model.reference.domain - name = base_model.reference.name - version = base_model.reference.version - if base_model.launchplan_ref is not None: - sdk_launch_plan = _launch_plan.SdkLaunchPlan.fetch(project, domain, name, version) - return cls(sdk_launch_plan=sdk_launch_plan) - elif base_model.sub_workflow_ref is not None: - # The workflow templates for sub-workflows should have been included in the original response - if base_model.reference in sub_workflows: - sw = sub_workflows[base_model.reference] - promoted = _workflow.SdkWorkflow.promote_from_model(sw, sub_workflows=sub_workflows, tasks=tasks) - return cls(sdk_workflow=promoted) - - # If not found for some reason, fetch it from Admin again. - # The reason there is a warning here but not for tasks is because sub-workflows should always be passed - # along. Ideally subworkflows are never even registered with Admin, so fetching from Admin ideally doesn't - # return anything. - _logging.warning( - "Your subworkflow with id {} is not included in the promote call.".format(base_model.reference) - ) - sdk_workflow = _workflow.SdkWorkflow.fetch(project, domain, name, version) - return cls(sdk_workflow=sdk_workflow) - else: - raise _system_exceptions.FlyteSystemException( - "Bad workflow node model, neither subworkflow nor " "launchplan specified." - ) diff --git a/flytekit/common/exceptions/__init__.py b/flytekit/common/exceptions/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/flytekit/common/interface.py b/flytekit/common/interface.py deleted file mode 100644 index c8f5160e15..0000000000 --- a/flytekit/common/interface.py +++ /dev/null @@ -1,163 +0,0 @@ -import six as _six - -from flytekit.common import promise as _promise -from flytekit.common import sdk_bases as _sdk_bases -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.types import containers as _containers -from flytekit.common.types import helpers as _type_helpers -from flytekit.common.types import primitives as _primitives -from flytekit.models import interface as _interface_models -from flytekit.models import literals as _literal_models - - -class BindingData(_literal_models.BindingData, metaclass=_sdk_bases.ExtendedSdkType): - @staticmethod - def _has_sub_bindings(m): - """ - :param dict[Text,T] or list[T]: - :rtype: bool - """ - for v in _six.itervalues(m) if isinstance(m, dict) else m: - if isinstance(v, (list, dict)) and BindingData._has_sub_bindings(v): - return True - elif isinstance(v, (_promise.Input, _promise.NodeOutput)): - return True - return False - - @classmethod - def promote_from_model(cls, model): - """ - :param flytekit.models.literals.BindingData model: - :rtype: BindingData - """ - return cls( - scalar=model.scalar, - collection=model.collection, - promise=model.promise, - map=model.map, - ) - - @classmethod - def from_python_std(cls, literal_type, t_value, upstream_nodes=None): - """ - :param flytekit.models.types.LiteralType literal_type: - :param T t_value: - :param list[flytekit.common.nodes.SdkNode] upstream_nodes: [Optional] Keeps track of the nodes upstream, - if applicable. - :rtype: BindingData - """ - scalar = None - collection = None - promise = None - map = None - downstream_sdk_type = _type_helpers.get_sdk_type_from_literal_type(literal_type) - if isinstance(t_value, _promise.Input): - if not downstream_sdk_type.is_castable_from(t_value.sdk_type): - _user_exceptions.FlyteTypeException( - t_value.sdk_type, - downstream_sdk_type, - additional_msg="When binding workflow input: {}".format(t_value), - ) - promise = t_value.promise - elif isinstance(t_value, _promise.NodeOutput): - if not downstream_sdk_type.is_castable_from(t_value.sdk_type): - _user_exceptions.FlyteTypeException( - t_value.sdk_type, - downstream_sdk_type, - additional_msg="When binding node output: {}".format(t_value), - ) - promise = t_value - if upstream_nodes is not None: - upstream_nodes.append(t_value.sdk_node) - elif isinstance(t_value, list): - if not issubclass(downstream_sdk_type, _containers.ListImpl): - raise _user_exceptions.FlyteTypeException( - type(t_value), - downstream_sdk_type, - received_value=t_value, - additional_msg="Cannot bind a list to a non-list type.", - ) - collection = _literal_models.BindingDataCollection( - [ - BindingData.from_python_std( - downstream_sdk_type.sub_type.to_flyte_literal_type(), - v, - upstream_nodes=upstream_nodes, - ) - for v in t_value - ] - ) - elif isinstance(t_value, dict) and ( - not issubclass(downstream_sdk_type, _primitives.Generic) or BindingData._has_sub_bindings(t_value) - ): - # TODO: This behavior should be embedded in the type engine. Someone should be able to alter behavior of - # TODO: binding logic by injecting their own type engine. The same goes for the list check above. - raise NotImplementedError("TODO: Cannot use map bindings at the moment") - else: - sdk_value = downstream_sdk_type.from_python_std(t_value) - scalar = sdk_value.scalar - collection = sdk_value.collection - map = sdk_value.map - return cls(scalar=scalar, collection=collection, map=map, promise=promise) - - -class TypedInterface(_interface_models.TypedInterface, metaclass=_sdk_bases.ExtendedSdkType): - @classmethod - def promote_from_model(cls, model): - """ - :param flytekit.models.interface.TypedInterface model: - :rtype: TypedInterface - """ - return cls(model.inputs, model.outputs) - - def create_bindings_for_inputs(self, map_of_bindings): - """ - :param dict[Text, T] map_of_bindings: This can be scalar primitives, it can be node output references, - lists, etc.. - :rtype: (list[flytekit.models.literals.Binding], list[flytekit.common.nodes.SdkNode]) - :raises: flytekit.common.exceptions.user.FlyteAssertion - """ - binding_data = dict() - all_upstream_nodes = list() - for k in sorted(self.inputs): - var = self.inputs[k] - if k not in map_of_bindings: - raise _user_exceptions.FlyteAssertion("Input was not specified for: {} of type {}".format(k, var.type)) - - binding_data[k] = BindingData.from_python_std( - var.type, map_of_bindings[k], upstream_nodes=all_upstream_nodes - ) - - extra_inputs = set(binding_data.keys()) ^ set(map_of_bindings.keys()) - if len(extra_inputs) > 0: - raise _user_exceptions.FlyteAssertion( - "Too many inputs were specified for the interface. Extra inputs were: {}".format(extra_inputs) - ) - - seen_nodes = set() - min_upstream = list() - for n in all_upstream_nodes: - if n not in seen_nodes: - seen_nodes.add(n) - min_upstream.append(n) - - return ( - [_literal_models.Binding(k, bd) for k, bd in _six.iteritems(binding_data)], - min_upstream, - ) - - def __repr__(self): - return "({inputs}) -> ({outputs})".format( - inputs=", ".join( - [ - "{}: {}".format(k, _type_helpers.get_sdk_type_from_literal_type(v.type)) - for k, v in _six.iteritems(self.inputs) - ] - ), - outputs=", ".join( - [ - "{}: {}".format(k, _type_helpers.get_sdk_type_from_literal_type(v.type)) - for k, v in _six.iteritems(self.outputs) - ] - ), - ) diff --git a/flytekit/common/launch_plan.py b/flytekit/common/launch_plan.py deleted file mode 100644 index adb9ff979d..0000000000 --- a/flytekit/common/launch_plan.py +++ /dev/null @@ -1,498 +0,0 @@ -import datetime as _datetime -import logging as _logging -import uuid as _uuid - -import six as _six -from deprecated import deprecated as _deprecated - -from flytekit.common import interface as _interface -from flytekit.common import nodes as _nodes -from flytekit.common import promise as _promises -from flytekit.common import sdk_bases as _sdk_bases -from flytekit.common import workflow_execution as _workflow_execution -from flytekit.common.core import identifier as _identifier -from flytekit.common.exceptions import scopes as _exception_scopes -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.mixins import hash as _hash_mixin -from flytekit.common.mixins import launchable as _launchable_mixin -from flytekit.common.mixins import registerable as _registerable -from flytekit.common.types import helpers as _type_helpers -from flytekit.configuration import auth as _auth_config -from flytekit.configuration import sdk as _sdk_config -from flytekit.engines.flyte import engine as _flyte_engine -from flytekit.models import common as _common_models -from flytekit.models import execution as _execution_models -from flytekit.models import interface as _interface_models -from flytekit.models import launch_plan as _launch_plan_models -from flytekit.models import literals as _literal_models -from flytekit.models import schedule as _schedule_model -from flytekit.models.core import identifier as _identifier_model -from flytekit.models.core import workflow as _workflow_models - - -class SdkLaunchPlan( - _launchable_mixin.LaunchableEntity, - _registerable.HasDependencies, - _registerable.RegisterableEntity, - _launch_plan_models.LaunchPlanSpec, - metaclass=_sdk_bases.ExtendedSdkType, -): - def __init__(self, *args, **kwargs): - super(SdkLaunchPlan, self).__init__(*args, **kwargs) - # Set all the attributes we expect this class to have - self._id = None - - # The interface is not set explicitly unless fetched in an engine context - self._interface = None - - @classmethod - def promote_from_model(cls, model) -> "SdkLaunchPlan": - """ - :param flytekit.models.launch_plan.LaunchPlanSpec model: - :rtype: SdkLaunchPlan - """ - return cls( - workflow_id=_identifier.Identifier.promote_from_model(model.workflow_id), - default_inputs=_interface_models.ParameterMap( - { - k: _promises.Input.promote_from_model(v).rename_and_return_reference(k) - for k, v in _six.iteritems(model.default_inputs.parameters) - } - ), - fixed_inputs=model.fixed_inputs, - entity_metadata=model.entity_metadata, - labels=model.labels, - annotations=model.annotations, - auth_role=model.auth_role, - raw_output_data_config=model.raw_output_data_config, - max_parallelism=model.max_parallelism, - ) - - @_exception_scopes.system_entry_point - def register(self, project, domain, name, version): - """ - :param Text project: - :param Text domain: - :param Text name: - :param Text version: - """ - self.validate() - id_to_register = _identifier.Identifier( - _identifier_model.ResourceType.LAUNCH_PLAN, project, domain, name, version - ) - client = _flyte_engine.get_client() - try: - client.create_launch_plan(id_to_register, self) - except _user_exceptions.FlyteEntityAlreadyExistsException: - pass - - self._id = id_to_register - self._has_registered = True - return str(self.id) - - @classmethod - @_exception_scopes.system_entry_point - def fetch(cls, project, domain, name, version=None): - """ - This function uses the engine loader to call create a hydrated task from Admin. - :param Text project: - :param Text domain: - :param Text name: - :param Text version: [Optional] If not set, the SDK will fetch the active launch plan for the given project, - domain, and name. - :rtype: SdkLaunchPlan - """ - from flytekit.common import workflow as _workflow - - launch_plan_id = _identifier.Identifier( - _identifier_model.ResourceType.LAUNCH_PLAN, project, domain, name, version - ) - - if launch_plan_id.version: - lp = _flyte_engine.get_client().get_launch_plan(launch_plan_id) - else: - named_entity_id = _common_models.NamedEntityIdentifier( - launch_plan_id.project, launch_plan_id.domain, launch_plan_id.name - ) - lp = _flyte_engine.get_client().get_active_launch_plan(named_entity_id) - - sdk_lp = cls.promote_from_model(lp.spec) - sdk_lp._id = lp.id - - # TODO: Add a test for this, and this function as a whole - wf_id = sdk_lp.workflow_id - lp_wf = _workflow.SdkWorkflow.fetch(wf_id.project, wf_id.domain, wf_id.name, wf_id.version) - sdk_lp._interface = lp_wf.interface - sdk_lp._has_registered = True - return sdk_lp - - @_exception_scopes.system_entry_point - def serialize(self): - """ - Serializing a launch plan should produce an object similar to what the registration step produces, - in preparation for actual registration to Admin. - - :rtype: flyteidl.admin.launch_plan_pb2.LaunchPlan - """ - return _launch_plan_models.LaunchPlan( - id=self.id, - spec=self, - closure=_launch_plan_models.LaunchPlanClosure( - state=None, - expected_inputs=_interface_models.ParameterMap({}), - expected_outputs=_interface_models.VariableMap({}), - ), - ).to_flyte_idl() - - @property - def id(self): - """ - :rtype: flytekit.common.core.identifier.Identifier - """ - return self._id - - @property - def is_scheduled(self): - """ - :rtype: bool - """ - if self.entity_metadata.schedule.cron_expression: - return True - elif self.entity_metadata.schedule.rate and self.entity_metadata.schedule.rate.value: - return True - elif self.entity_metadata.schedule.cron_schedule and self.entity_metadata.schedule.cron_schedule.schedule: - return True - else: - return False - - @property - def auth_role(self): - """ - :rtype: flytekit.models.common.AuthRole - """ - fixed_auth = super(SdkLaunchPlan, self).auth_role - if fixed_auth is not None and ( - fixed_auth.assumable_iam_role is not None or fixed_auth.kubernetes_service_account is not None - ): - return fixed_auth - - assumable_iam_role = _auth_config.ASSUMABLE_IAM_ROLE.get() - kubernetes_service_account = _auth_config.KUBERNETES_SERVICE_ACCOUNT.get() - - if not assumable_iam_role and _sdk_config.ROLE.get() is not None: - _logging.warning( - "Using deprecated `role` from config. Please update your config to use `assumable_iam_role` instead" - ) - assumable_iam_role = _sdk_config.ROLE.get() - return _common_models.AuthRole( - assumable_iam_role=assumable_iam_role, - kubernetes_service_account=kubernetes_service_account, - ) - - @property - def workflow_id(self): - """ - :rtype: flytekit.common.core.identifier.Identifier - """ - return self._workflow_id - - @property - def interface(self): - """ - The interface is not technically part of the admin.LaunchPlanSpec in the IDL, however the workflow ID is, and - from the workflow ID, fetch will fill in the interface. This is nice because then you can __call__ the= - object and get a node. - :rtype: flytekit.common.interface.TypedInterface - """ - return self._interface - - @property - def resource_type(self): - """ - Integer from _identifier.ResourceType enum - :rtype: int - """ - return _identifier_model.ResourceType.LAUNCH_PLAN - - @property - def entity_type_text(self): - """ - :rtype: Text - """ - return "Launch Plan" - - @property - def raw_output_data_config(self): - """ - :rtype: flytekit.models.common.RawOutputDataConfig - """ - raw_output_data_config = super(SdkLaunchPlan, self).raw_output_data_config - if raw_output_data_config is not None and raw_output_data_config.output_location_prefix != "": - return raw_output_data_config - - # If it was not set explicitly then let's use the value found in the configuration. - return _common_models.RawOutputDataConfig(_auth_config.RAW_OUTPUT_DATA_PREFIX.get()) - - @_exception_scopes.system_entry_point - def validate(self): - # TODO: Validate workflow is satisfied - pass - - @_exception_scopes.system_entry_point - def update(self, state): - """ - :param int state: Enum value from flytekit.models.launch_plan.LaunchPlanState - """ - if not self.id: - raise _user_exceptions.FlyteAssertion( - "Failed to update launch plan because the launch plan's ID is not set. Please call register to fetch " - "or register the identifier first" - ) - return _flyte_engine.get_client().update_launch_plan(self.id, state) - - def _python_std_input_map_to_literal_map(self, inputs): - """ - :param dict[Text,Any] inputs: A dictionary of Python standard inputs that will be type-checked and compiled - to a LiteralMap - :rtype: flytekit.models.literals.LiteralMap - """ - return _type_helpers.pack_python_std_map_to_literal_map( - inputs, - {k: user_input.sdk_type for k, user_input in _six.iteritems(self.default_inputs.parameters) if k in inputs}, - ) - - @_deprecated(reason="Use launch_with_literals instead", version="0.9.0") - def execute_with_literals( - self, - project, - domain, - literal_inputs, - name=None, - notification_overrides=None, - label_overrides=None, - annotation_overrides=None, - ): - """ - Deprecated. - """ - return self.launch_with_literals( - project, - domain, - literal_inputs, - name, - notification_overrides, - label_overrides, - annotation_overrides, - ) - - @_exception_scopes.system_entry_point - def launch_with_literals( - self, - project, - domain, - literal_inputs, - name=None, - notification_overrides=None, - label_overrides=None, - annotation_overrides=None, - auth_role=None, - ): - """ - Executes the launch plan and returns the execution identifier. This version of execution is meant for when - you already have a LiteralMap of inputs. - - :param Text project: - :param Text domain: - :param flytekit.models.literals.LiteralMap literal_inputs: Inputs to the execution. - :param Text name: [Optional] If specified, an execution will be created with this name. Note: the name must - be unique within the context of the project and domain. - :param list[flytekit.common.notifications.Notification] notification_overrides: [Optional] If specified, these - are the notifications that will be honored for this execution. An empty list signals to disable all - notifications. - :param flytekit.models.common.Labels label_overrides: - :param flytekit.models.common.Annotations annotation_overrides: - :rtype: flytekit.common.workflow_execution.SdkWorkflowExecution - :param flytekit.models.common.AuthRole auth_role: - """ - # Kubernetes requires names starting with an alphabet for some resources. - name = name or "f" + _uuid.uuid4().hex[:19] - disable_all = notification_overrides == [] - if disable_all: - notification_overrides = None - else: - notification_overrides = _execution_models.NotificationList(notification_overrides or []) - disable_all = None - - client = _flyte_engine.get_client() - try: - exec_id = client.create_execution( - project, - domain, - name, - _execution_models.ExecutionSpec( - self.id, - _execution_models.ExecutionMetadata( - _execution_models.ExecutionMetadata.ExecutionMode.MANUAL, - "sdk", # TODO: get principle - 0, # TODO: Detect nesting - ), - notifications=notification_overrides, - disable_all=disable_all, - labels=label_overrides, - annotations=annotation_overrides, - auth_role=auth_role, - ), - literal_inputs, - ) - except _user_exceptions.FlyteEntityAlreadyExistsException: - exec_id = _identifier.WorkflowExecutionIdentifier(project, domain, name) - execution = client.get_execution(exec_id) - return _workflow_execution.SdkWorkflowExecution.promote_from_model(execution) - - @_exception_scopes.system_entry_point - def __call__(self, *args, **input_map): - """ - :param list[T] args: Do not specify. Kwargs only are supported for this function. - :param dict[Text,T] input_map: Map of inputs. Can be statically defined or OutputReference links. - :rtype: flytekit.common.nodes.SdkNode - """ - if len(args) > 0: - raise _user_exceptions.FlyteAssertion( - "When adding a launchplan as a node in a workflow, all inputs must be specified with kwargs only. We " - "detected {} positional args.".format(len(args)) - ) - - # Take the default values from the launch plan - default_inputs = {k: v.sdk_default for k, v in _six.iteritems(self.default_inputs.parameters) if not v.required} - default_inputs.update(input_map) - - bindings, upstream_nodes = self.interface.create_bindings_for_inputs(default_inputs) - - return _nodes.SdkNode( - id=None, - metadata=_workflow_models.NodeMetadata("", _datetime.timedelta(), _literal_models.RetryStrategy(0)), - bindings=sorted(bindings, key=lambda b: b.var), - upstream_nodes=upstream_nodes, - sdk_launch_plan=self, - ) - - def __repr__(self): - """ - :rtype: Text - """ - return "SdkLaunchPlan(ID: {} Interface: {} WF ID: {})".format(self.id, self.interface, self.workflow_id) - - -# The difference between this and the SdkLaunchPlan class is that this runnable class is supposed to only be used for -# launch plans loaded alongside the current Python interpreter. -class SdkRunnableLaunchPlan(_hash_mixin.HashOnReferenceMixin, SdkLaunchPlan): - def __init__( - self, - sdk_workflow, - default_inputs=None, - fixed_inputs=None, - role=None, - schedule=None, - notifications=None, - labels=None, - annotations=None, - auth_role=None, - raw_output_data_config=None, - ): - """ - :param flytekit.common.local_workflow.SdkRunnableWorkflow sdk_workflow: - :param dict[Text,flytekit.common.promise.Input] default_inputs: - :param dict[Text,Any] fixed_inputs: These inputs will be fixed and not need to be set when executing this - launch plan. - :param Text role: Deprecated. IAM role to execute this launch plan with. - :param flytekit.models.schedule.Schedule: Schedule to apply to this workflow. - :param list[flytekit.models.common.Notification]: List of notifications to apply to this launch plan. - :param flytekit.models.common.Labels labels: Any custom kubernetes labels to apply to workflows executed by this - launch plan. - :param flytekit.models.common.Annotations annotations: Any custom kubernetes annotations to apply to workflows - executed by this launch plan. - Any custom kubernetes annotations to apply to workflows executed by this launch plan. - :param flytekit.models.common.Authrole auth_role: The auth method with which to execute the workflow. - :param flytekit.models.common.RawOutputDataConfig raw_output_data_config: Config for offloading data - """ - if role and auth_role: - raise ValueError("Cannot set both role and auth. Role is deprecated, use auth instead.") - - fixed_inputs = fixed_inputs or {} - default_inputs = default_inputs or {} - - if role: - auth_role = _common_models.AuthRole(assumable_iam_role=role) - - # The constructor for SdkLaunchPlan sets the id to None anyways so we don't bother passing in an ID. The ID - # should be set in one of three places, - # 1) When the object is registered (in the code above) - # 2) By the dynamic task code after this runnable object has already been __call__'ed. The SdkNode produced - # maintains a link to this object and will set the ID according to the configuration variables present. - # 3) When SdkLaunchPlan.fetch() is run - super(SdkRunnableLaunchPlan, self).__init__( - None, - _launch_plan_models.LaunchPlanMetadata( - schedule=schedule or _schedule_model.Schedule(""), - notifications=notifications or [], - ), - _interface_models.ParameterMap(default_inputs), - _type_helpers.pack_python_std_map_to_literal_map( - fixed_inputs, - { - k: _type_helpers.get_sdk_type_from_literal_type(var.type) - for k, var in _six.iteritems(sdk_workflow.interface.inputs) - if k in fixed_inputs - }, - ), - labels or _common_models.Labels({}), - annotations or _common_models.Annotations({}), - auth_role, - raw_output_data_config or _common_models.RawOutputDataConfig(""), - ) - self._interface = _interface.TypedInterface( - {k: v.var for k, v in _six.iteritems(default_inputs)}, - sdk_workflow.interface.outputs, - ) - self._upstream_entities = {sdk_workflow} - self._sdk_workflow = sdk_workflow - - @classmethod - def from_flyte_idl(cls, _): - raise _user_exceptions.FlyteAssertion( - "An SdkRunnableLaunchPlan must be created from a reference to local Python code only." - ) - - @classmethod - def promote_from_model(cls, model): - raise _user_exceptions.FlyteAssertion( - "An SdkRunnableLaunchPlan must be created from a reference to local Python code only." - ) - - @classmethod - @_exception_scopes.system_entry_point - def fetch(cls, project, domain, name, version=None): - """ - This function uses the engine loader to call create a hydrated task from Admin. - :param Text project: - :param Text domain: - :param Text name: - :param Text version: - :rtype: SdkRunnableLaunchPlan - """ - raise _user_exceptions.FlyteAssertion( - "An SdkRunnableLaunchPlan must be created from a reference to local Python code only." - ) - - @property - def workflow_id(self): - """ - :rtype: flytekit.common.core.identifier.Identifier - """ - return self._sdk_workflow.id - - def __repr__(self): - """ - :rtype: Text - """ - return "SdkRunnableLaunchPlan(ID: {} Interface: {} WF ID: {})".format(self.id, self.interface, self.workflow_id) diff --git a/flytekit/common/local_workflow.py b/flytekit/common/local_workflow.py deleted file mode 100644 index eb2067578f..0000000000 --- a/flytekit/common/local_workflow.py +++ /dev/null @@ -1,388 +0,0 @@ -import uuid as _uuid -from typing import Any, Dict, List - -import six as _six -from six.moves import queue as _queue - -from flytekit.common import interface as _interface -from flytekit.common import launch_plan as _launch_plan -from flytekit.common import nodes as _nodes -from flytekit.common import promise as _promise -from flytekit.common.core import identifier as _identifier -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.types import helpers as _type_helpers -from flytekit.common.workflow import SdkWorkflow -from flytekit.configuration import internal as _internal_config -from flytekit.models import common as _common_models -from flytekit.models import interface as _interface_models -from flytekit.models import literals as _literal_models -from flytekit.models import schedule as _schedule_models -from flytekit.models.core import identifier as _identifier_model -from flytekit.models.core import workflow as _workflow_models - - -# Local-only wrapper around binding data and variables. Note that the Output object used by the end user is a yet -# another layer on top of this. -class Output(object): - def __init__(self, name, value, sdk_type=None, help=None): - """ - :param Text name: - :param T value: - :param U sdk_type: If specified, the value provided must cast to this type. Normally should be an instance of - flytekit.common.types.base_sdk_types.FlyteSdkType. But could also be something like: - - list[flytekit.common.types.base_sdk_types.FlyteSdkType], - dict[flytekit.common.types.base_sdk_types.FlyteSdkType,flytekit.common.types.base_sdk_types.FlyteSdkType], - (flytekit.common.types.base_sdk_types.FlyteSdkType, flytekit.common.types.base_sdk_types.FlyteSdkType, ...) - """ - if sdk_type is None: - # This syntax didn't work for some reason: sdk_type = sdk_type or Output._infer_type(value) - sdk_type = Output._infer_type(value) - sdk_type = _type_helpers.python_std_to_sdk_type(sdk_type) - - self._binding_data = _interface.BindingData.from_python_std(sdk_type.to_flyte_literal_type(), value) - self._var = _interface_models.Variable(sdk_type.to_flyte_literal_type(), help or "") - self._name = name - - def rename_and_return_reference(self, new_name): - self._name = new_name - return self - - @staticmethod - def _infer_type(value): - # TODO: Infer types - raise NotImplementedError( - "Currently the SDK cannot infer a workflow output type, so please use the type kwarg " - "when instantiating an output." - ) - - @property - def name(self): - """ - :rtype: Text - """ - return self._name - - @property - def binding_data(self): - """ - :rtype: flytekit.models.literals.BindingData - """ - return self._binding_data - - @property - def var(self): - """ - :rtype: flytekit.models.interface.Variable - """ - return self._var - - -class SdkRunnableWorkflow(SdkWorkflow): - """ - Wrapper class for workflows defined using Python, written in a Flyte workflow repo. This class is misnamed. It is - more appropriately called PythonWorkflow. The reason we are calling it SdkRunnableWorkflow instead is merely in - keeping with the established convention in other parts of this codebase. We will likely change this naming scheme - entirely before a 1.0 release. - - Being "runnable" or not "runnable" is not a distinction we care to make at this point. Do not read into it, - pretend it's not there. The purpose of this class is merely to differentiate between - i) A workflow object, as created by a user's workflow code, using the @workflow_class decorator for instance. If - you have one of these classes, it means you have the actual Python code available in the Python process you - are running. - ii) The SdkWorkflow object, which represents a workflow as retrieved from Flyte Admin. Anyone with access - to Admin, can instantiate an SdkWorkflow object by 'fetch'ing it. You don't need to have any of - the actual code checked out. SdkWorkflow's are effectively then the control plane model of a workflow, - represented as a Python object. - """ - - def __init__( - self, - inputs: List[_promise.Input], - nodes: List[_nodes.SdkNode], - interface, - output_bindings, - id=None, - metadata=None, - metadata_defaults=None, - disable_default_launch_plan=False, - ): - """ - :param list[flytekit.common.nodes.SdkNode] nodes: - :param flytekit.models.interface.TypedInterface interface: Defines a strongly typed interface for the - Workflow (inputs, outputs). This can include some optional parameters. - :param list[flytekit.models.literals.Binding] output_bindings: A list of output bindings that specify how to construct - workflow outputs. Bindings can pull node outputs or specify literals. All workflow outputs specified in - the interface field must be bound - in order for the workflow to be validated. A workflow has an implicit dependency on all of its nodes - to execute successfully in order to bind final outputs. - :param flytekit.models.core.identifier.Identifier id: This is an autogenerated id by the system. The id is - globally unique across Flyte. - :param WorkflowMetadata metadata: This contains information on how to run the workflow. - :param flytekit.models.core.workflow.WorkflowMetadataDefaults metadata_defaults: Defaults to be passed - to nodes contained within workflow. - :param bool disable_default_launch_plan: Determines whether to create a default launch plan for the workflow. - """ - # Save the promise.Input objects for future use. - self._user_inputs = inputs - - # Set optional settings - id = ( - id - if id is not None - else _identifier.Identifier( - _identifier_model.ResourceType.WORKFLOW, - _internal_config.PROJECT.get(), - _internal_config.DOMAIN.get(), - _uuid.uuid4().hex, - _internal_config.VERSION.get(), - ) - ) - metadata = metadata if metadata is not None else _workflow_models.WorkflowMetadata() - metadata_defaults = ( - metadata_defaults if metadata_defaults is not None else _workflow_models.WorkflowMetadataDefaults() - ) - - super(SdkRunnableWorkflow, self).__init__( - nodes=nodes, - interface=interface, - output_bindings=output_bindings, - id=id, - metadata=metadata, - metadata_defaults=metadata_defaults, - ) - - # Set this last as it's set in constructor - self._upstream_entities = set(n.executable_sdk_object for n in nodes) - self._should_create_default_launch_plan = not disable_default_launch_plan - - @property - def should_create_default_launch_plan(self): - """ - Determines whether registration flow should create a default launch plan for this workflow or not. - :rtype: bool - """ - return self._should_create_default_launch_plan - - def __call__(self, *args, **input_map): - # Take the default values from the Inputs - compiled_inputs = {v.name: v.sdk_default for v in self.user_inputs if not v.sdk_required} - compiled_inputs.update(input_map) - - return super().__call__(*args, **compiled_inputs) - - @classmethod - def construct_from_class_definition( - cls, - inputs: List[_promise.Input], - outputs: List[Output], - nodes: List[_nodes.SdkNode], - metadata: _workflow_models.WorkflowMetadata = None, - metadata_defaults: _workflow_models.WorkflowMetadataDefaults = None, - disable_default_launch_plan: bool = False, - ) -> "SdkRunnableWorkflow": - """ - This constructor is here to provide backwards-compatibility for class-defined Workflows - - :param list[flytekit.common.promise.Input] inputs: - :param list[Output] outputs: - :param list[flytekit.common.nodes.SdkNode] nodes: - :param WorkflowMetadata metadata: This contains information on how to run the workflow. - :param flytekit.models.core.workflow.WorkflowMetadataDefaults metadata_defaults: Defaults to be passed - to nodes contained within workflow. - :param bool disable_default_launch_plan: Determines whether to create a default launch plan for the workflow or not. - - :rtype: SdkRunnableWorkflow - """ - for n in nodes: - for upstream in n.upstream_nodes: - if upstream.id is None: - raise _user_exceptions.FlyteAssertion( - "Some nodes contained in the workflow were not found in the workflow description. Please " - "ensure all nodes are either assigned to attributes within the class or an element in a " - "list, dict, or tuple which is stored as an attribute in the class." - ) - - id = _identifier.Identifier( - _identifier_model.ResourceType.WORKFLOW, - _internal_config.PROJECT.get(), - _internal_config.DOMAIN.get(), - _uuid.uuid4().hex, - _internal_config.VERSION.get(), - ) - interface = _interface.TypedInterface({v.name: v.var for v in inputs}, {v.name: v.var for v in outputs}) - - output_bindings = [_literal_models.Binding(v.name, v.binding_data) for v in outputs] - - return cls( - inputs=inputs, - nodes=nodes, - interface=interface, - output_bindings=output_bindings, - id=id, - metadata=metadata, - metadata_defaults=metadata_defaults, - disable_default_launch_plan=disable_default_launch_plan, - ) - - @property - def id(self): - return self._id - - @id.setter - def id(self, new_id): - self._id = new_id - - @property - def user_inputs(self) -> List[_promise.Input]: - """ - :rtype: list[flytekit.common.promise.Input] - """ - return self._user_inputs - - def create_launch_plan( - self, - default_inputs: Dict[str, _promise.Input] = None, - fixed_inputs: Dict[str, Any] = None, - schedule: _schedule_models.Schedule = None, - role: str = None, - notifications: List[_common_models.Notification] = None, - labels: _common_models.Labels = None, - annotations: _common_models.Annotations = None, - assumable_iam_role: str = None, - kubernetes_service_account: str = None, - raw_output_data_prefix: str = None, - ): - """ - This method will create a launch plan object that can execute this workflow. - :param dict[Text,flytekit.common.promise.Input] default_inputs: - :param dict[Text,T] fixed_inputs: - :param flytekit.models.schedule.Schedule schedule: A schedule on which to execute this launch plan. - :param Text role: Deprecated. Use assumable_iam_role instead. - :param list[flytekit.models.common.Notification] notifications: A list of notifications to enact by default for - this launch plan. - :param flytekit.models.common.Labels labels: - :param flytekit.models.common.Annotations annotations: - :param cls: This parameter can be used by users to define an extension of a launch plan to instantiate. The - class provided should be a subclass of flytekit.common.launch_plan.SdkLaunchPlan. - :param Text assumable_iam_role: The IAM role to execute the workflow with. - :param Text kubernetes_service_account: The kubernetes service account to execute the workflow with. - :param Text raw_output_data_prefix: Bucket for offloaded data - :rtype: flytekit.common.launch_plan.SdkRunnableLaunchPlan - """ - # TODO: Actually ensure the parameters conform. - if role and (assumable_iam_role or kubernetes_service_account): - raise ValueError("Cannot set both role and auth. Role is deprecated, use auth instead.") - fixed_inputs = fixed_inputs or {} - merged_default_inputs = {v.name: v for v in self.user_inputs if v.name not in fixed_inputs} - merged_default_inputs.update(default_inputs or {}) - - if role: - assumable_iam_role = role # For backwards compatibility - auth_role = _common_models.AuthRole( - assumable_iam_role=assumable_iam_role, - kubernetes_service_account=kubernetes_service_account, - ) - - raw_output_config = _common_models.RawOutputDataConfig(raw_output_data_prefix or "") - - return _launch_plan.SdkRunnableLaunchPlan( - sdk_workflow=self, - default_inputs={ - k: user_input.rename_and_return_reference(k) for k, user_input in _six.iteritems(merged_default_inputs) - }, - fixed_inputs=fixed_inputs, - schedule=schedule, - notifications=notifications, - labels=labels, - annotations=annotations, - auth_role=auth_role, - raw_output_data_config=raw_output_config, - ) - - -def build_sdk_workflow_from_metaclass(metaclass, on_failure=None, disable_default_launch_plan=False, cls=None): - """ - :param T metaclass: This is the user-defined workflow class, prior to decoration. - :param on_failure flytekit.models.core.workflow.WorkflowMetadata.OnFailurePolicy: [Optional] The execution policy - when the workflow detects a failure. - :param bool disable_default_launch_plan: Determines whether to create a default launch plan for the workflow or not. - :param cls: This is the class that will be instantiated from the inputs, outputs, and nodes. This will be used - by users extending the base Flyte programming model. If set, it must be a subclass of PythonWorkflow. - - :rtype: SdkRunnableWorkflow - """ - inputs, outputs, nodes = _discover_workflow_components(metaclass) - metadata = _workflow_models.WorkflowMetadata(on_failure=on_failure if on_failure else None) - - return (cls or SdkRunnableWorkflow).construct_from_class_definition( - inputs=[i for i in sorted(inputs, key=lambda x: x.name)], - outputs=[o for o in sorted(outputs, key=lambda x: x.name)], - nodes=[n for n in sorted(nodes, key=lambda x: x.id)], - metadata=metadata, - disable_default_launch_plan=disable_default_launch_plan, - ) - - -def _discover_workflow_components(workflow_class): - """ - This task iterates over the attributes of a user-defined class in order to return a list of inputs, outputs and - nodes. - :param class workflow_class: User-defined class with task instances as attributes. - :rtype: (list[flytekit.common.promise.Input], list[Output], list[flytekit.common.nodes.SdkNode]) - """ - - inputs = [] - outputs = [] - nodes = [] - - to_visit_objs = _queue.Queue() - top_level_attributes = set() - for attribute_name in dir(workflow_class): - to_visit_objs.put((attribute_name, getattr(workflow_class, attribute_name))) - top_level_attributes.add(attribute_name) - - # For all task instances defined within the workflow, bind them to this specific workflow and hook-up to the - # engine (when available) - visited_obj_ids = set() - while not to_visit_objs.empty(): - attribute_name, current_obj = to_visit_objs.get() - - current_obj_id = id(current_obj) - if current_obj_id in visited_obj_ids: - continue - visited_obj_ids.add(current_obj_id) - - if isinstance(current_obj, _nodes.SdkNode): - # TODO: If an attribute name is on the form node_name[index], the resulting - # node name might not be correct. - nodes.append(current_obj.assign_id_and_return(attribute_name)) - elif isinstance(current_obj, _promise.Input): - if attribute_name is None or attribute_name not in top_level_attributes: - raise _user_exceptions.FlyteValueException( - attribute_name, - "Detected workflow input specified outside of top level.", - ) - inputs.append(current_obj.rename_and_return_reference(attribute_name)) - elif isinstance(current_obj, Output): - if attribute_name is None or attribute_name not in top_level_attributes: - raise _user_exceptions.FlyteValueException( - attribute_name, - "Detected workflow output specified outside of top level.", - ) - outputs.append(current_obj.rename_and_return_reference(attribute_name)) - elif isinstance(current_obj, list) or isinstance(current_obj, set) or isinstance(current_obj, tuple): - for idx, value in enumerate(current_obj): - to_visit_objs.put((_assign_indexed_attribute_name(attribute_name, idx), value)) - elif isinstance(current_obj, dict): - # Visit dictionary keys. - for key in current_obj.keys(): - to_visit_objs.put((_assign_indexed_attribute_name(attribute_name, key), key)) - # Visit dictionary values. - for key, value in _six.iteritems(current_obj): - to_visit_objs.put((_assign_indexed_attribute_name(attribute_name, key), value)) - return inputs, outputs, nodes - - -def _assign_indexed_attribute_name(attribute_name, index): - return "{}[{}]".format(attribute_name, index) diff --git a/flytekit/common/mixins/__init__.py b/flytekit/common/mixins/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/flytekit/common/mixins/artifact.py b/flytekit/common/mixins/artifact.py deleted file mode 100644 index 41e5ebb0b7..0000000000 --- a/flytekit/common/mixins/artifact.py +++ /dev/null @@ -1,80 +0,0 @@ -import abc as _abc -import datetime as _datetime -import time as _time - -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.models import common as _common_models - - -class ExecutionArtifact(object, metaclass=_common_models.FlyteABCMeta): - @_abc.abstractproperty - def inputs(self): - """ - Returns the inputs to the execution in the standard Python format as dictated by the type engine. - :rtype: dict[Text, T] - """ - pass - - @_abc.abstractproperty - def outputs(self): - """ - Returns the outputs to the execution in the standard Python format as dictated by the type engine. If the - execution ended in error or the execution is in progress, an exception will be raised. - :rtype: dict[Text, T] - """ - pass - - @_abc.abstractproperty - def error(self): - """ - If execution is in progress, raise an exception. Otherwise, return None if no error was present upon - reaching completion. - :rtype: flytekit.models.core.execution.ExecutionError or None - """ - pass - - @_abc.abstractproperty - def is_complete(self): - """ - Dictates whether or not the execution is complete. - :rtype: bool - """ - pass - - @_abc.abstractmethod - def sync(self): - """ - Syncs the state of the underlying execution artifact with the state observed by the platform. - :rtype: None - """ - pass - - @_abc.abstractmethod - def _sync_closure(self): - """ - Syncs the closure of the underlying execution artifact with the state observed by the platform. - :rtype: None - """ - pass - - def wait_for_completion(self, timeout=None, poll_interval=None): - """ - :param datetime.timedelta timeout: Amount of time to wait until the execution has completed before timing - out. If not set or set to None, this method will wait for infinite. - :param datetime.timedelta poll_interval: Duration to wait between polling for a completion update. - :rtype: None - """ - poll_interval = poll_interval or _datetime.timedelta(seconds=30) - if timeout is None: - time_to_give_up = _datetime.datetime.max - else: - time_to_give_up = _datetime.datetime.utcnow() + timeout - - self._sync_closure() - while _datetime.datetime.utcnow() < time_to_give_up: - if self.is_complete: - self.sync() - return - _time.sleep(poll_interval.total_seconds()) - self._sync_closure() - raise _user_exceptions.FlyteTimeout("Execution {} did not complete before timeout.".format(self)) diff --git a/flytekit/common/mixins/launchable.py b/flytekit/common/mixins/launchable.py deleted file mode 100644 index 110ba663af..0000000000 --- a/flytekit/common/mixins/launchable.py +++ /dev/null @@ -1,129 +0,0 @@ -import abc as _abc - -from deprecated import deprecated as _deprecated - - -class LaunchableEntity(object, metaclass=_abc.ABCMeta): - def launch( - self, - project, - domain, - inputs=None, - name=None, - notification_overrides=None, - label_overrides=None, - annotation_overrides=None, - auth_role=None, - ): - """ - Creates a remote execution from the entity and returns the execution identifier. - This version of launch is meant for when inputs are specified as Python native types/structures. - - :param Text project: - :param Text domain: - :param dict[Text, Any] inputs: A dictionary of Python standard inputs that will be type-checked, then compiled - to a LiteralMap. - :param Text name: [Optional] If specified, an execution will be created with this name. Note: the name must - be unique within the context of the project and domain. - :param list[flytekit.common.notifications.Notification] notification_overrides: [Optional] If specified, these - are the notifications that will be honored for this execution. An empty list signals to disable all - notifications. - :param flytekit.models.common.Labels label_overrides: - :param flytekit.models.common.Annotations annotation_overrides: - :param flytekit.models.common.AuthRole auth_role: - :rtype: T - - """ - return self.launch_with_literals( - project, - domain, - self._python_std_input_map_to_literal_map(inputs or {}), - name=name, - notification_overrides=notification_overrides, - label_overrides=label_overrides, - annotation_overrides=annotation_overrides, - auth_role=auth_role, - ) - - @_deprecated(reason="Use launch instead", version="0.9.0") - def execute( - self, - project, - domain, - inputs=None, - name=None, - notification_overrides=None, - label_overrides=None, - annotation_overrides=None, - ): - """ - Deprecated. - """ - return self.launch( - project, - domain, - inputs=inputs, - name=name, - notification_overrides=notification_overrides, - label_overrides=label_overrides, - annotation_overrides=annotation_overrides, - ) - - @_abc.abstractmethod - def _python_std_input_map_to_literal_map(self, inputs): - pass - - @_abc.abstractmethod - def launch_with_literals( - self, - project, - domain, - literal_inputs, - name=None, - notification_overrides=None, - label_overrides=None, - annotation_overrides=None, - auth_role=None, - ): - """ - Executes the entity and returns the execution identifier. This version of execution is meant for when - you already have a LiteralMap of inputs. - - :param Text project: - :param Text domain: - :param flytekit.models.literals.LiteralMap literal_inputs: Inputs to the execution. - :param Text name: [Optional] If specified, an execution will be created with this name. Note: the name must - be unique within the context of the project and domain. - :param list[flytekit.common.notifications.Notification] notification_overrides: [Optional] If specified, these - are the notifications that will be honored for this execution. An empty list signals to disable all - notifications. - :param flytekit.models.common.Labels label_overrides: - :param flytekit.models.common.Annotations annotation_overrides: - :param flytekit.models.common.AuthRole auth_role: - :rtype: flytekit.models.core.identifier.WorkflowExecutionIdentifier: - """ - pass - - @_deprecated(reason="Use launch_with_literals instead", version="0.9.0") - def execute_with_literals( - self, - project, - domain, - literal_inputs, - name=None, - notification_overrides=None, - label_overrides=None, - annotation_overrides=None, - ): - """ - Deprecated. - """ - return self.launch_with_literals( - project, - domain, - literal_inputs, - name, - notification_overrides, - label_overrides, - annotation_overrides, - ) diff --git a/flytekit/common/mixins/registerable.py b/flytekit/common/mixins/registerable.py deleted file mode 100644 index 64fba47ef0..0000000000 --- a/flytekit/common/mixins/registerable.py +++ /dev/null @@ -1,195 +0,0 @@ -import abc as _abc -import importlib as _importlib -import inspect as _inspect -import logging as _logging -from typing import Set - -from flytekit.common import sdk_bases as _sdk_bases -from flytekit.common import utils as _utils -from flytekit.common.exceptions import system as _system_exceptions - - -class _InstanceTracker(_sdk_bases.ExtendedSdkType): - """ - This is either genius or terrible. Some of our tools iterate over modules and try to find Flyte entities - (Tasks, Workflows, Launch Plans) and then register them. However, if a task is imported via a command like: - - from package.module import some_task - - It is possible we will find a task reference twice, but then how do we know where it was defined? Ideally, we would - like to only register a task once and do so with the name where it is defined. This metaclass allows us to do this - by inspecting the call stack when __call__ is called on the metaclass (thus instantiating an object). - """ - - @staticmethod - def _find_instance_module(): - frame = _inspect.currentframe() - while frame: - if frame.f_code.co_name == "": - return frame.f_globals["__name__"] - frame = frame.f_back - return None - - def __call__(cls, *args, **kwargs): - o = super(_InstanceTracker, cls).__call__(*args, **kwargs) - o._instantiated_in = _InstanceTracker._find_instance_module() - return o - - -class FlyteEntity(object, metaclass=_sdk_bases.ExtendedSdkType): - @property - @_abc.abstractmethod - def resource_type(self): - """ - Integer from _identifier.ResourceType enum - :rtype: int - """ - pass - - @property - @_abc.abstractmethod - def entity_type_text(self): - """ - TODO: Rename to resource type text - :rtype: Text - """ - pass - - -class TrackableEntity(FlyteEntity, metaclass=_InstanceTracker): - def __init__(self, *args, **kwargs): - self._platform_valid_name = None - super(TrackableEntity, self).__init__(*args, **kwargs) - - @property - def instantiated_in(self): - """ - If found, we try to specify the module where the task was first instantiated. - :rtype: Optional[Text] - """ - return self._instantiated_in # Set in metaclass - - @property - def has_valid_name(self): - """ - :rtype: bool - """ - return self._platform_valid_name is not None and self._platform_valid_name != "" - - @property - def platform_valid_name(self): - """ - :rtype: Text - """ - return self._platform_valid_name - - def assign_name(self, name): - self._platform_valid_name = name - - def auto_assign_name(self): - """ - This function is a bit of trickster Python code that goes hand in hand with the _InstanceTracker metaclass - defined above. Thanks @matthewphsmith for this bit of ingenuity. - - For instance, if a user has code that looks like this: - - from some.other.module import wf - my_launch_plan = wf.create_launch_plan() - - @dynamic_task - def sample_task(wf_params): - yield my_launch_plan() - - This code means that we should have a launch plan with a name ending in "my_launch_plan", since that is the - name of the variable that the created launch plan gets assigned to. That is also the name that the launch plan - would be registered with. - - However, when the create_launch_plan() function runs, the Python interpreter has no idea where the created - object will be assigned to. It has no idea that the output of the create_launch_plan call is to be paired up - with a variable named "my_launch_plan". This function basically does this after the fact. Leveraging the - _instantiated_in field provided by the _InstanceTracker class above, this code will re-import the - module (ie Python file) that the object is in. Since it's already loaded, it's just retrieved from memory. - It then scans all objects in the module, and when an object match is found, it knows it's found the right - variable name. - - Just to drive the point home, this function is mostly needed for Launch Plans. Assuming that user code has: - - @python_task - def some_task() - - When Flytekit calls the module loader and loads the task, the name of the task is the name of the function - itself. It's known at time of creation. In contrast, when - - xyz = SomeWorkflow.create_launch_plan() - - is called, the name of the launch plan isn't known until after creation, it's not "SomeWorkflow", it's "xyz" - """ - _logging.debug("Running name auto assign") - m = _importlib.import_module(self.instantiated_in) - - for k in dir(m): - try: - if getattr(m, k) is self: - self._platform_valid_name = _utils.fqdn(m.__name__, k, entity_type=self.resource_type) - _logging.debug("Auto-assigning name to {}".format(self._platform_valid_name)) - return - except ValueError as err: - # Empty pandas dataframes behave weirdly here such that calling `m.df` raises: - # ValueError: The truth value of a {type(self).__name__} is ambiguous. Use a.empty, a.bool(), a.item(), - # a.any() or a.all() - # Since dataframes aren't registrable entities to begin with we swallow any errors they raise and - # continue looping through m. - _logging.warning("Caught ValueError {} while attempting to auto-assign name".format(err)) - pass - - _logging.error("Could not auto-assign name") - raise _system_exceptions.FlyteSystemException("Error looking for object while auto-assigning name.") - - -class RegisterableEntity(TrackableEntity): - def __init__(self, *args, **kwargs): - self._has_registered = False - super(RegisterableEntity, self).__init__(*args, **kwargs) - - @_abc.abstractmethod - def register(self, project, domain, name, version): - """ - :param Text project: The project in which to register this task. - :param Text domain: The domain in which to register this task. - :param Text name: The name to give this task. - :param Text version: The version in which to register this task. - """ - pass - - @_abc.abstractmethod - def serialize(self): - """ - Registerable entities also are required to be serialized. This allows flytekit to separate serialization from - the network call to Admin (mostly at least, if a Launch Plan is fetched for instance as part of another - workflow, it will still hit Admin). - """ - pass - - @property - def has_registered(self) -> bool: - return self._has_registered - - -class HasDependencies(object): - """ - This interface is meant to describe Flyte entities that can have upstream dependencies. For instance, currently a - launch plan depends on the underlying workflow, and a workflow is dependent on its tasks, and other launch plans, - and subworkflows. - """ - - def __init__(self, *args, **kwargs): - self._upstream_entities = set() - super(HasDependencies, self).__init__(*args, **kwargs) - - @property - def upstream_entities(self) -> Set[RegisterableEntity]: - """ - Task, workflow, and launch plan that need to be registered in advance of this workflow. - :rtype: set[RegisterableEntity] - """ - return self._upstream_entities diff --git a/flytekit/common/nodes.py b/flytekit/common/nodes.py deleted file mode 100644 index fb68de2bc2..0000000000 --- a/flytekit/common/nodes.py +++ /dev/null @@ -1,463 +0,0 @@ -import abc as _abc -import logging as _logging -import os as _os - -import six as _six -from flyteidl.core import literals_pb2 as _literals_pb2 -from sortedcontainers import SortedDict as _SortedDict - -from flytekit.clients.helpers import iterate_task_executions as _iterate_task_executions -from flytekit.common import component_nodes as _component_nodes -from flytekit.common import constants as _constants -from flytekit.common import promise as _promise -from flytekit.common import sdk_bases as _sdk_bases -from flytekit.common import utils as _common_utils -from flytekit.common.exceptions import scopes as _exception_scopes -from flytekit.common.exceptions import system as _system_exceptions -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.mixins import artifact as _artifact_mixin -from flytekit.common.mixins import hash as _hash_mixin -from flytekit.common.tasks import executions as _task_executions -from flytekit.common.types import helpers as _type_helpers -from flytekit.common.utils import _dnsify -from flytekit.engines.flyte import engine as _flyte_engine -from flytekit.interfaces.data import data_proxy as _data_proxy -from flytekit.models import common as _common_models -from flytekit.models import literals as _literal_models -from flytekit.models import node_execution as _node_execution_models -from flytekit.models.core import execution as _execution_models -from flytekit.models.core import workflow as _workflow_model - - -class ParameterMapper(_SortedDict, metaclass=_common_models.FlyteABCMeta): - """ - This abstract class provides functionality to reference specific inputs and outputs for a task instance. This - allows for syntax such as: - - my_task_instance.inputs.my_input - - And is especially useful for linking tasks together via outputs -> inputs in workflow definitions: - - my_second_task_instance(input=my_task_instances.outputs.my_output) - - Attributes: - Dynamically discovered. Only the keys for inputs/outputs can be referenced. - - Example: - - .. code-block:: python - - @inputs(a=Types.Integer) - @outputs(b=Types.String) - @python_task(version='1') - def my_task(wf_params, a, b): - pass - - input_link = my_task.inputs.a # Success! - output_link = my_tasks.outputs.b # Success! - - input_link = my_task.inputs.c # Attribute not found exception! - output_link = my_task.outputs.d # Attribute not found exception! - - """ - - def __init__(self, type_map, node): - """ - :param dict[Text, flytekit.models.interface.Variable] type_map: - :param SdkNode node: - """ - super(ParameterMapper, self).__init__() - for key, var in _six.iteritems(type_map): - self[key] = self._return_mapping_object(node, _type_helpers.get_sdk_type_from_literal_type(var.type), key) - self._initialized = True - - def __getattr__(self, key): - if key == "iteritems" and hasattr(super(ParameterMapper, self), "items"): - return super(ParameterMapper, self).items - if hasattr(super(ParameterMapper, self), key): - return getattr(super(ParameterMapper, self), key) - if key not in self: - raise _user_exceptions.FlyteAssertion("{} doesn't exist.".format(key)) - return self[key] - - def __setattr__(self, key, value): - if "_initialized" in self.__dict__: - raise _user_exceptions.FlyteAssertion("Parameters are immutable.") - else: - super(ParameterMapper, self).__setattr__(key, value) - - @_abc.abstractmethod - def _return_mapping_object(self, sdk_node, sdk_type, name): - """ - :param flytekit.common.nodes.Node sdk_node: - :param flytekit.common.types.FlyteSdkType sdk_type: - :param Text name: - """ - pass - - -class OutputParameterMapper(ParameterMapper): - """ - This subclass of ParameterMapper is used to represent outputs for a given node. - """ - - def _return_mapping_object(self, sdk_node, sdk_type, name): - """ - :param flytekit.common.nodes.Node sdk_node: - :param flytekit.common.types.FlyteSdkType sdk_type: - :param Text name: - """ - return _promise.NodeOutput(sdk_node, sdk_type, name) - - -class SdkNode(_hash_mixin.HashOnReferenceMixin, _workflow_model.Node, metaclass=_sdk_bases.ExtendedSdkType): - def __init__( - self, - id, - upstream_nodes, - bindings, - metadata, - sdk_task=None, - sdk_workflow=None, - sdk_launch_plan=None, - sdk_branch=None, - parameter_mapping=True, - ): - """ - :param Text id: A workflow-level unique identifier that identifies this node in the workflow. "inputs" and - "outputs" are reserved node ids that cannot be used by other nodes. - :param flytekit.models.core.workflow.NodeMetadata metadata: Extra metadata about the node. - :param list[flytekit.models.literals.Binding] bindings: Specifies how to bind the underlying - interface's inputs. All required inputs specified in the underlying interface must be fulfilled. - :param list[SdkNode] upstream_nodes: Specifies execution dependencies for this node ensuring it will - only get scheduled to run after all its upstream nodes have completed. This node will have - an implicit dependency on any node that appears in inputs field. - :param flytekit.common.tasks.task.SdkTask sdk_task: The task to execute in this - node. - :param flytekit.common.workflow.SdkWorkflow sdk_workflow: The workflow to execute in this node. - :param flytekit.common.launch_plan.SdkLaunchPlan sdk_launch_plan: The launch plan to execute in this - node. - :param TODO sdk_branch: TODO - """ - non_none_entities = [ - entity for entity in [sdk_workflow, sdk_branch, sdk_launch_plan, sdk_task] if entity is not None - ] - if len(non_none_entities) != 1: - raise _user_exceptions.FlyteAssertion( - "An SDK node must have one underlying entity specified at once. Received the following " - "entities: {}".format(non_none_entities) - ) - - workflow_node = None - if sdk_workflow is not None: - workflow_node = _component_nodes.SdkWorkflowNode(sdk_workflow=sdk_workflow) - elif sdk_launch_plan is not None: - workflow_node = _component_nodes.SdkWorkflowNode(sdk_launch_plan=sdk_launch_plan) - - # TODO: this calls the constructor which means it will set all the upstream node ids to None if at the time of - # this instantiation, the upstream nodes have not had their nodes assigned yet. - super(SdkNode, self).__init__( - id=_dnsify(id) if id else None, - metadata=metadata, - inputs=bindings, - upstream_node_ids=[n.id for n in upstream_nodes], - output_aliases=[], # TODO: Are aliases a thing in SDK nodes - task_node=_component_nodes.SdkTaskNode(sdk_task) if sdk_task else None, - workflow_node=workflow_node, - branch_node=sdk_branch, - ) - self._upstream = upstream_nodes - self._executable_sdk_object = sdk_task or sdk_workflow or sdk_launch_plan - if parameter_mapping: - if not sdk_branch: - self._outputs = OutputParameterMapper(self._executable_sdk_object.interface.outputs, self) - else: - self._outputs = None - - @property - def executable_sdk_object(self): - return self._executable_sdk_object - - @classmethod - def promote_from_model(cls, model, sub_workflows, tasks): - """ - :param flytekit.models.core.workflow.Node model: - :param dict[flytekit.models.core.identifier.Identifier, flytekit.models.core.workflow.WorkflowTemplate] - sub_workflows: - :param dict[flytekit.models.core.identifier.Identifier, flytekit.models.task.TaskTemplate] tasks: If specified, - these task templates will be passed to the SdkTaskNode promote_from_model call, and used - instead of fetching from Admin. - :rtype: SdkNode - """ - id = model.id - # This should never be called - if id == _constants.START_NODE_ID or id == _constants.END_NODE_ID: - _logging.warning("Should not call promote from model on a start node or end node {}".format(model)) - return None - - sdk_task_node, sdk_workflow_node = None, None - if model.task_node is not None: - sdk_task_node = _component_nodes.SdkTaskNode.promote_from_model(model.task_node, tasks) - elif model.workflow_node is not None: - sdk_workflow_node = _component_nodes.SdkWorkflowNode.promote_from_model( - model.workflow_node, sub_workflows, tasks - ) - else: - raise _system_exceptions.FlyteSystemException("Bad Node model, neither task nor workflow detected") - - # When WorkflowTemplate models (containing node models) are returned by Admin, they've been compiled with a - # start node. In order to make the promoted SdkWorkflow look the same, we strip the start-node text back out. - for i in model.inputs: - if i.binding.promise is not None and i.binding.promise.node_id == _constants.START_NODE_ID: - i.binding.promise._node_id = _constants.GLOBAL_INPUT_NODE_ID - - if sdk_task_node is not None: - return cls( - id=id, - upstream_nodes=[], # set downstream, model doesn't contain this information - bindings=model.inputs, - metadata=model.metadata, - sdk_task=sdk_task_node.sdk_task, - ) - elif sdk_workflow_node is not None: - if sdk_workflow_node.sdk_workflow is not None: - return cls( - id=id, - upstream_nodes=[], # set downstream, model doesn't contain this information - bindings=model.inputs, - metadata=model.metadata, - sdk_workflow=sdk_workflow_node.sdk_workflow, - ) - elif sdk_workflow_node.sdk_launch_plan is not None: - return cls( - id=id, - upstream_nodes=[], # set downstream, model doesn't contain this information - bindings=model.inputs, - metadata=model.metadata, - sdk_launch_plan=sdk_workflow_node.sdk_launch_plan, - ) - else: - raise _system_exceptions.FlyteSystemException( - "Bad SdkWorkflowNode model, both lp and workflow are None" - ) - else: - raise _system_exceptions.FlyteSystemException("Bad SdkNode model, both task and workflow nodes are empty") - - @property - def upstream_nodes(self): - """ - :rtype: list[SdkNode] - """ - return self._upstream - - @property - def upstream_node_ids(self): - """ - :rtype: list[Text] - """ - return [n.id for n in sorted(self.upstream_nodes, key=lambda x: x.id)] - - @property - def outputs(self): - """ - :rtype: dict[Text, flytekit.common.promise.NodeOutput] - """ - return self._outputs - - def assign_id_and_return(self, id): - """ - :param Text id: - :rtype: None - """ - if self.id: - raise _user_exceptions.FlyteAssertion( - "Error assigning ID: {} because {} is already assigned. Has this node been assigned to another " - "workflow already?".format(id, self) - ) - self._id = _dnsify(id) if id else None - self._metadata._name = id - return self - - def with_overrides(self, *args, **kwargs): - # TODO: Implement overrides - raise NotImplementedError("Overrides are not supported in Flyte yet.") - - @_exception_scopes.system_entry_point - def __lshift__(self, other): - """ - Add a node upstream of this node without necessarily mapping outputs -> inputs. - :param Node other: node to place upstream - """ - if hash(other) not in set(hash(n) for n in self.upstream_nodes): - self._upstream.append(other) - return other - - @_exception_scopes.system_entry_point - def __rshift__(self, other): - """ - Add a node downstream of this node without necessarily mapping outputs -> inputs. - - :param Node other: node to place downstream - """ - if hash(self) not in set(hash(n) for n in other.upstream_nodes): - other.upstream_nodes.append(self) - return other - - def __repr__(self): - """ - :rtype: Text - """ - return "Node(ID: {} Executable: {})".format(self.id, self._executable_sdk_object) - - -class SdkNodeExecution( - _node_execution_models.NodeExecution, _artifact_mixin.ExecutionArtifact, metaclass=_sdk_bases.ExtendedSdkType -): - def __init__(self, *args, **kwargs): - super(SdkNodeExecution, self).__init__(*args, **kwargs) - self._task_executions = None - self._workflow_executions = None - self._inputs = None - self._outputs = None - - @property - def task_executions(self): - """ - Returns the underlying task executions in order of try attempt. - :rtype: list[flytekit.common.tasks.executions.SdkTaskExecution] - """ - return self._task_executions or [] - - @property - def workflow_executions(self): - """ - Returns the underlying workflow executions in order of try attempt. - :rtype: list[flytekit.common.workflow_execution.SdkWorkflowExecution] - """ - return self._workflow_executions or [] - - @property - def executions(self): - """ - Returns a list of generic execution artifacts. - :rtype: list[flytekit.common.mixins.artifact.ExecutionArtifact] - """ - return self.task_executions or self.workflow_executions or [] - - @property - def inputs(self): - """ - Returns the inputs to the execution in the standard Python format as dictated by the type engine. - :rtype: dict[Text, T] - """ - if self._inputs is None: - client = _flyte_engine.get_client() - execution_data = client.get_node_execution_data(self.id) - - # Inputs are returned inline unless they are too big, in which case a url blob pointing to them is returned. - if bool(execution_data.full_inputs.literals): - input_map = execution_data.full_inputs - elif execution_data.inputs.bytes > 0: - with _common_utils.AutoDeletingTempDir() as t: - tmp_name = _os.path.join(t.name, "inputs.pb") - _data_proxy.Data.get_data(execution_data.inputs.url, tmp_name) - input_map = _literal_models.LiteralMap.from_flyte_idl( - _common_utils.load_proto_from_file(_literals_pb2.LiteralMap, tmp_name) - ) - else: - input_map = _literal_models.LiteralMap({}) - - self._inputs = _type_helpers.unpack_literal_map_to_sdk_python_std(input_map) - return self._inputs - - @property - def outputs(self): - """ - Returns the outputs to the execution in the standard Python format as dictated by the type engine. If the - execution ended in error or the execution is in progress, an exception will be raised. - :rtype: dict[Text, T] - """ - if not self.is_complete: - raise _user_exceptions.FlyteAssertion( - "Please what until the node execution has completed before requesting the outputs." - ) - if self.error: - raise _user_exceptions.FlyteAssertion("Outputs could not be found because the execution ended in failure.") - - if self._outputs is None: - client = _flyte_engine.get_client() - execution_data = client.get_node_execution_data(self.id) - - # Outputs are returned inline unless they are too big, in which case a url blob pointing to them is returned. - if bool(execution_data.full_outputs.literals): - output_map = execution_data.full_outputs - - elif execution_data.outputs.bytes > 0: - with _common_utils.AutoDeletingTempDir() as t: - tmp_name = _os.path.join(t.name, "outputs.pb") - _data_proxy.Data.get_data(execution_data.outputs.url, tmp_name) - output_map = _literal_models.LiteralMap.from_flyte_idl( - _common_utils.load_proto_from_file(_literals_pb2.LiteralMap, tmp_name) - ) - else: - output_map = _literal_models.LiteralMap({}) - - self._outputs = _type_helpers.unpack_literal_map_to_sdk_python_std(output_map) - return self._outputs - - @property - def error(self): - """ - If execution is in progress, raise an exception. Otherwise, return None if no error was present upon - reaching completion. - :rtype: flytekit.models.core.execution.ExecutionError or None - """ - if not self.is_complete: - raise _user_exceptions.FlyteAssertion( - "Please wait until the node execution has completed before requesting error information." - ) - return self.closure.error - - @property - def is_complete(self): - """ - Dictates whether or not the execution is complete. - :rtype: bool - """ - return self.closure.phase in { - _execution_models.NodeExecutionPhase.ABORTED, - _execution_models.NodeExecutionPhase.FAILED, - _execution_models.NodeExecutionPhase.SKIPPED, - _execution_models.NodeExecutionPhase.SUCCEEDED, - _execution_models.NodeExecutionPhase.TIMED_OUT, - } - - @classmethod - def promote_from_model(cls, base_model): - """ - :param _node_execution_models.NodeExecution base_model: - :rtype: SdkNodeExecution - """ - return cls( - closure=base_model.closure, id=base_model.id, input_uri=base_model.input_uri, metadata=base_model.metadata - ) - - def sync(self): - """ - Syncs the state of this object with that held by the platform. - :rtype: None - """ - if not self.is_complete or self.task_executions is not None: - client = _flyte_engine.get_client() - self._closure = client.get_node_execution(self.id).closure - task_executions = list(_iterate_task_executions(client, self.id)) - self._task_executions = [_task_executions.SdkTaskExecution.promote_from_model(te) for te in task_executions] - # TODO: Sub-workflows too once implemented - - def _sync_closure(self): - """ - Syncs the closure of the underlying execution artifact with the state observed by the platform. - :rtype: None - """ - client = _flyte_engine.get_client() - self._closure = client.get_node_execution(self.id).closure diff --git a/flytekit/common/notifications.py b/flytekit/common/notifications.py deleted file mode 100644 index 09b1d11358..0000000000 --- a/flytekit/common/notifications.py +++ /dev/null @@ -1,101 +0,0 @@ -from flytekit.common import sdk_bases as _sdk_bases -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.models import common as _common_model -from flytekit.models.core import execution as _execution_model - - -class Notification(_common_model.Notification, metaclass=_sdk_bases.ExtendedSdkType): - - VALID_PHASES = { - _execution_model.WorkflowExecutionPhase.ABORTED, - _execution_model.WorkflowExecutionPhase.FAILED, - _execution_model.WorkflowExecutionPhase.SUCCEEDED, - _execution_model.WorkflowExecutionPhase.TIMED_OUT, - } - - def __init__(self, phases, email=None, pager_duty=None, slack=None): - """ - :param list[int] phases: A required list of phases for which to fire the event. Events can only be fired for - terminal phases. Phases should be as defined in: flytekit.models.core.execution.WorkflowExecutionPhase - """ - self._validate_phases(phases) - super(Notification, self).__init__(phases, email=email, pager_duty=pager_duty, slack=slack) - - def _validate_phases(self, phases): - """ - :param list[int] phases: - """ - if len(phases) == 0: - raise _user_exceptions.FlyteAssertion("You must specify at least one phase for a notification.") - for phase in phases: - if phase not in self.VALID_PHASES: - raise _user_exceptions.FlyteValueException( - phase, - self.VALID_PHASES, - additional_message="Notifications can only be specified on terminal states.", - ) - - @classmethod - def from_flyte_idl(cls, p): - """ - :param flyteidl.admin.common_pb2.Notification p: - :rtype: Notification - """ - if p.HasField("email"): - return cls(p.phases, p.email.recipients_email) - elif p.HasField("pager_duty"): - return cls(p.phases, p.pager_duty.recipients_email) - else: - return cls(p.phases, p.slack.recipients_email) - - -class PagerDuty(Notification): - def __init__(self, phases, recipients_email): - """ - :param list[Text] recipients_email: A required non-empty list of recipients for the notification. - """ - super(PagerDuty, self).__init__(phases, pager_duty=_common_model.PagerDutyNotification(recipients_email)) - - @classmethod - def promote_from_model(cls, base_model): - """ - :param flytekit.models.common.Notification base_model: - :rtype: Notification - """ - return cls(base_model.phases, base_model.pager_duty.recipients_email) - - -class Email(Notification): - def __init__(self, phases, recipients_email): - """ - :param list[Text] recipients_email: A required non-empty list of recipients for the notification. - :param list[int] phases: A required list of phases for which to fire the event. Events can only be fired for - terminal phases. Phases should be as defined in: flytekit.models.core.execution.WorkflowExecutionPhase - """ - super(Email, self).__init__(phases, email=_common_model.EmailNotification(recipients_email)) - - @classmethod - def promote_from_model(cls, base_model): - """ - :param flytekit.models.common.Notification base_model: - :rtype: Notification - """ - return cls(base_model.phases, base_model.email.recipients_email) - - -class Slack(Notification): - def __init__(self, phases, recipients_email): - """ - :param list[Text] recipients_email: A required non-empty list of recipients for the notification. - :param list[int] phases: A required list of phases for which to fire the event. Events can only be fired for - terminal phases. Phases should be as defined in: flytekit.models.core.execution.WorkflowExecutionPhase - """ - super(Slack, self).__init__(phases, slack=_common_model.SlackNotification(recipients_email)) - - @classmethod - def promote_from_model(cls, base_model): - """ - :param flytekit.models.common.Notification base_model: - :rtype: Notification - """ - return cls(base_model.phases, base_model.slack.recipients_email) diff --git a/flytekit/common/promise.py b/flytekit/common/promise.py deleted file mode 100644 index 79352dd638..0000000000 --- a/flytekit/common/promise.py +++ /dev/null @@ -1,169 +0,0 @@ -from flytekit.common import constants as _constants -from flytekit.common import sdk_bases as _sdk_bases -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.types import helpers as _type_helpers -from flytekit.models import interface as _interface_models -from flytekit.models import types as _type_models - - -class Input(_interface_models.Parameter, metaclass=_sdk_bases.ExtendedSdkType): - def __init__(self, name, sdk_type, help=None, **kwargs): - """ - :param Text name: - :param flytekit.common.types.base_sdk_types.FlyteSdkType sdk_type: This is the SDK type necessary to create an - input to this workflow. - :param Text help: An optional help string to describe the input to users. - :param bool required: If set to True, default must be None - :param T default: If this is not a required input, the value will default to this value. - """ - param_default = None - if "required" not in kwargs and "default" not in kwargs: - # Neither required or default is set so assume required - required = True - default = None - elif kwargs.get("required", False) and "default" in kwargs: - # Required cannot be set to True and have a default specified - raise _user_exceptions.FlyteAssertion("Default cannot be set when required is True") - elif "default" in kwargs: - # If default is specified, then required must be false and the value is whatever is specified - required = None - default = kwargs["default"] - param_default = sdk_type.from_python_std(default) - else: - # If no default is set, but required is set, then the behavior is determined by required == True or False - default = None - required = kwargs["required"] - if not required: - # If required == False, we assume default to be None - param_default = sdk_type.from_python_std(default) - required = None - - self._sdk_required = required or False - self._sdk_default = default - self._help = help - self._sdk_type = sdk_type - self._promise = _type_models.OutputReference(_constants.GLOBAL_INPUT_NODE_ID, name) - self._name = name - super(Input, self).__init__( - _interface_models.Variable(type=sdk_type.to_flyte_literal_type(), description=help or ""), - required=required, - default=param_default, - ) - - def rename_and_return_reference(self, new_name): - self._promise._var = new_name - return self - - @property - def name(self): - """ - :rtype: Text - """ - return self._promise.var - - @property - def promise(self): - """ - :rtype: flytekit.models.types.OutputReference - """ - return self._promise - - @property - def sdk_required(self): - """ - :rtype: bool - """ - return self._sdk_required - - @property - def sdk_default(self): - """ - :rtype: T - """ - return self._sdk_default - - @property - def help(self): - """ - :rtype: Text - """ - return self._help - - @property - def sdk_type(self): - """ - :rtype: flytekit.common.types.base_sdk_types.FlyteSdkType - """ - return self._sdk_type - - def __repr__(self): - return "Input({}, {}, required={}, help={})".format(self.name, self.sdk_type, self.required, self.help) - - @classmethod - def promote_from_model(cls, model): - """ - :param flytekit.models.interface.Parameter model: - :rtype: Input - """ - sdk_type = _type_helpers.get_sdk_type_from_literal_type(model.var.type) - - if model.default is not None: - default_value = sdk_type.from_flyte_idl(model.default.to_flyte_idl()).to_python_std() - return cls( - "", - sdk_type, - help=model.var.description, - required=False, - default=default_value, - ) - else: - return cls("", sdk_type, help=model.var.description, required=True) - - -class NodeOutput(_type_models.OutputReference, metaclass=_sdk_bases.ExtendedSdkType): - def __init__(self, sdk_node, sdk_type, var): - """ - :param sdk_node: - :param sdk_type: deprecated in mypy flytekit. - :param var: - """ - self._node = sdk_node - self._type = sdk_type - super(NodeOutput, self).__init__(self._node.id, var) - - @property - def node_id(self): - """ - Override the underlying node_id property to refer to SdkNode. - :rtype: Text - """ - return self.sdk_node.id - - @classmethod - def promote_from_model(cls, model): - """ - :param flytekit.models.types.OutputReference model: - :rtype: NodeOutput - """ - raise _user_exceptions.FlyteAssertion( - "A NodeOutput cannot be promoted from a protobuf because it must be " - "contextualized by an existing SdkNode." - ) - - @property - def sdk_node(self): - """ - :rtype: flytekit.common.nodes.SdkNode - """ - return self._node - - @property - def sdk_type(self): - """ - :rtype: flytekit.common.types.base_sdk_types.FlyteSdkType - """ - return self._type - - def __repr__(self): - s = f"NodeOutput({self.sdk_node if self.sdk_node.id is not None else None}:{self.var})" - return s diff --git a/flytekit/common/schedules.py b/flytekit/common/schedules.py deleted file mode 100644 index d6e71c6f83..0000000000 --- a/flytekit/common/schedules.py +++ /dev/null @@ -1,195 +0,0 @@ -import datetime as _datetime -import re as _re - -import croniter as _croniter - -from flytekit.common import sdk_bases as _sdk_bases -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.models import schedule as _schedule_models - - -class _ExtendedSchedule(_schedule_models.Schedule): - @classmethod - def from_flyte_idl(cls, proto): - """ - :param flyteidl.admin.schedule_pb2.Schedule proto: - :rtype: _ExtendedSchedule - """ - return cls.promote_from_model(_schedule_models.Schedule.from_flyte_idl(proto)) - - -class CronSchedule(_ExtendedSchedule, metaclass=_sdk_bases.ExtendedSdkType): - _VALID_CRON_ALIASES = [ - "hourly", - "hours", - "@hourly", - "daily", - "days", - "@daily", - "weekly", - "weeks", - "@weekly", - "monthly", - "months", - "@monthly", - "annually", - "@annually", - "yearly", - "years", - "@yearly", - ] - - # Not a perfect regex but good enough and simple to reason about - _OFFSET_PATTERN = _re.compile("([-+]?)P([-+0-9YMWD]+)?(T([-+0-9HMS.,]+)?)?") - - def __init__(self, cron_expression=None, schedule=None, offset=None, kickoff_time_input_arg=None): - """ - :param Text cron_expression: - :param Text schedule: - :param Text offset: - :param Text kickoff_time_input_arg: - """ - if cron_expression is None and schedule is None: - raise _user_exceptions.FlyteAssertion("Either `cron_expression` or `schedule` should be specified.") - - if cron_expression is not None and offset is not None: - raise _user_exceptions.FlyteAssertion("Only `schedule` is supported when specifying `offset`.") - - if cron_expression is not None: - CronSchedule._validate_expression(cron_expression) - - if schedule is not None: - CronSchedule._validate_schedule(schedule) - - if offset is not None: - CronSchedule._validate_offset(offset) - - super(CronSchedule, self).__init__( - kickoff_time_input_arg, - cron_expression=cron_expression, - cron_schedule=_schedule_models.Schedule.CronSchedule(schedule, offset) if schedule is not None else None, - ) - - @staticmethod - def _validate_expression(cron_expression): - """ - Ensures that the set value is a valid cron string. We use the format used in Cloudwatch and the best - explanation can be found here: - https://docs.aws.amazon.com/AmazonCloudWatch/latest/events/ScheduledEvents.html#CronExpressions - :param Text cron_expression: cron expression - """ - # We use the croniter lib to validate our cron expression. Since on the admin side we use Cloudwatch, - # we have a couple checks in order to line up Cloudwatch with Croniter. - tokens = cron_expression.split() - if len(tokens) != 6: - raise _user_exceptions.FlyteAssertion( - "Cron expression is invalid. A cron expression must have 6 fields. Cron expressions are in the " - "format of: `minute hour day-of-month month day-of-week year`. " - "Use `schedule` for 5 fields cron expression. Received: `{}`".format(cron_expression) - ) - - if tokens[2] != "?" and tokens[4] != "?": - raise _user_exceptions.FlyteAssertion( - "Scheduled string is invalid. A cron expression must have a '?' for either day-of-month or " - "day-of-week. Please specify '?' for one of those fields. Cron expressions are in the format of: " - "minute hour day-of-month month day-of-week year.\n\n" - "For more information: " - "https://docs.aws.amazon.com/AmazonCloudWatch/latest/events/ScheduledEvents.html#CronExpressions" - ) - - try: - # Cut to 5 fields and just assume year field is good because croniter treats the 6th field as seconds. - # TODO: Parse this field ourselves and check - _croniter.croniter(" ".join(cron_expression.replace("?", "*").split()[:5])) - except Exception: - raise _user_exceptions.FlyteAssertion( - "Scheduled string is invalid. The cron expression was found to be invalid." - " Provided cron expr: {}".format(cron_expression) - ) - - @staticmethod - def _validate_schedule(schedule): - if schedule.lower() not in CronSchedule._VALID_CRON_ALIASES: - try: - _croniter.croniter(schedule) - except Exception: - raise _user_exceptions.FlyteAssertion( - "Schedule is invalid. It must be set to either a cron alias or valid cron expression." - " Provided schedule: {}".format(schedule) - ) - - @staticmethod - def _validate_offset(offset): - if CronSchedule._OFFSET_PATTERN.fullmatch(offset) is None: - raise _user_exceptions.FlyteAssertion( - "Offset is invalid. It must be an ISO 8601 duration. Provided offset: {}".format(offset) - ) - - @classmethod - def promote_from_model(cls, base_model): - """ - :param flytekit.models.schedule.Schedule base_model: - :rtype: CronSchedule - """ - return cls( - cron_expression=base_model.cron_expression, - schedule=base_model.cron_schedule.schedule if base_model.cron_schedule is not None else None, - offset=base_model.cron_schedule.offset if base_model.cron_schedule is not None else None, - kickoff_time_input_arg=base_model.kickoff_time_input_arg, - ) - - -class FixedRate(_ExtendedSchedule, metaclass=_sdk_bases.ExtendedSdkType): - def __init__(self, duration, kickoff_time_input_arg=None): - """ - :param datetime.timedelta duration: - :param Text kickoff_time_input_arg: - """ - super(FixedRate, self).__init__(kickoff_time_input_arg, rate=self._translate_duration(duration)) - - @staticmethod - def _translate_duration(duration): - """ - :param datetime.timedelta duration: timedelta between runs - :rtype: flytekit.models.schedule.Schedule.FixedRate - """ - _SECONDS_TO_MINUTES = 60 - _SECONDS_TO_HOURS = _SECONDS_TO_MINUTES * 60 - _SECONDS_TO_DAYS = _SECONDS_TO_HOURS * 24 - - if duration.microseconds != 0 or duration.seconds % _SECONDS_TO_MINUTES != 0: - raise _user_exceptions.FlyteAssertion( - "Granularity of less than a minute is not supported for FixedRate schedules. Received: {}".format( - duration - ) - ) - elif int(duration.total_seconds()) % _SECONDS_TO_DAYS == 0: - return _schedule_models.Schedule.FixedRate( - int(duration.total_seconds() / _SECONDS_TO_DAYS), - _schedule_models.Schedule.FixedRateUnit.DAY, - ) - elif int(duration.total_seconds()) % _SECONDS_TO_HOURS == 0: - return _schedule_models.Schedule.FixedRate( - int(duration.total_seconds() / _SECONDS_TO_HOURS), - _schedule_models.Schedule.FixedRateUnit.HOUR, - ) - else: - return _schedule_models.Schedule.FixedRate( - int(duration.total_seconds() / _SECONDS_TO_MINUTES), - _schedule_models.Schedule.FixedRateUnit.MINUTE, - ) - - @classmethod - def promote_from_model(cls, base_model): - """ - :param flytekit.models.schedule.Schedule base_model: - :rtype: FixedRate - """ - if base_model.rate.unit == _schedule_models.Schedule.FixedRateUnit.DAY: - duration = _datetime.timedelta(days=base_model.rate.value) - elif base_model.rate.unit == _schedule_models.Schedule.FixedRateUnit.HOUR: - duration = _datetime.timedelta(hours=base_model.rate.value) - else: - duration = _datetime.timedelta(minutes=base_model.rate.value) - - return cls(duration, kickoff_time_input_arg=base_model.kickoff_time_input_arg) diff --git a/flytekit/common/sdk_bases.py b/flytekit/common/sdk_bases.py deleted file mode 100644 index d082302343..0000000000 --- a/flytekit/common/sdk_bases.py +++ /dev/null @@ -1,22 +0,0 @@ -import abc as _abc - -from flytekit.models import common as _common - - -class ExtendedSdkType(_common.FlyteType, metaclass=_common.FlyteABCMeta): - """ - Abstract class that all SDK objects must inherit from. This provides the ability to promote a data model object - into an actionable object. - """ - - @_abc.abstractmethod - def promote_from_model(cls, base_model): - """ - :param flytekit.models.common.FlyteIdlEntity base_model: - :rtype: ExtendedSdkType - """ - pass - - def from_flyte_idl(cls, pb2_object): - base_model = super(ExtendedSdkType, cls).from_flyte_idl(pb2_object) - return cls.promote_from_model(base_model) diff --git a/flytekit/common/tasks/__init__.py b/flytekit/common/tasks/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/flytekit/common/tasks/executions.py b/flytekit/common/tasks/executions.py deleted file mode 100644 index d87c558c09..0000000000 --- a/flytekit/common/tasks/executions.py +++ /dev/null @@ -1,153 +0,0 @@ -import os as _os - -import six as _six -from flyteidl.core import literals_pb2 as _literals_pb2 - -from flytekit.clients.helpers import iterate_node_executions as _iterate_node_executions -from flytekit.common import sdk_bases as _sdk_bases -from flytekit.common import utils as _common_utils -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.mixins import artifact as _artifact_mixin -from flytekit.common.types import helpers as _type_helpers -from flytekit.engines.flyte import engine as _flyte_engine -from flytekit.interfaces.data import data_proxy as _data_proxy -from flytekit.models import literals as _literal_models -from flytekit.models.admin import task_execution as _task_execution_model -from flytekit.models.core import execution as _execution_models - - -class SdkTaskExecution( - _task_execution_model.TaskExecution, _artifact_mixin.ExecutionArtifact, metaclass=_sdk_bases.ExtendedSdkType -): - def __init__(self, *args, **kwargs): - super(SdkTaskExecution, self).__init__(*args, **kwargs) - self._inputs = None - self._outputs = None - - @property - def is_complete(self): - """ - Dictates whether or not the execution is complete. - :rtype: bool - """ - return self.closure.phase in { - _execution_models.TaskExecutionPhase.ABORTED, - _execution_models.TaskExecutionPhase.FAILED, - _execution_models.TaskExecutionPhase.SUCCEEDED, - } - - @property - def inputs(self): - """ - Returns the inputs of the task execution in the standard Python format that is produced by - the type engine. - :rtype: dict[Text, T] - """ - if self._inputs is None: - client = _flyte_engine.get_client() - execution_data = client.get_task_execution_data(self.id) - - # Inputs are returned inline unless they are too big, in which case a url blob pointing to them is returned. - if bool(execution_data.full_inputs.literals): - input_map = execution_data.full_inputs - elif execution_data.inputs.bytes > 0: - with _common_utils.AutoDeletingTempDir() as t: - tmp_name = _os.path.join(t.name, "inputs.pb") - _data_proxy.Data.get_data(execution_data.inputs.url, tmp_name) - input_map = _literal_models.LiteralMap.from_flyte_idl( - _common_utils.load_proto_from_file(_literals_pb2.LiteralMap, tmp_name) - ) - else: - input_map = _literal_models.LiteralMap({}) - - self._inputs = _type_helpers.unpack_literal_map_to_sdk_python_std(input_map) - return self._inputs - - @property - def outputs(self): - """ - Returns the outputs of the task execution, if available, in the standard Python format that is produced by - the type engine. If not available, perhaps due to execution being in progress or an error being produced, - this will raise an exception. - :rtype: dict[Text, T] - """ - if not self.is_complete: - raise _user_exceptions.FlyteAssertion( - "Please what until the task execution has completed before requesting the outputs." - ) - if self.error: - raise _user_exceptions.FlyteAssertion("Outputs could not be found because the execution ended in failure.") - - if self._outputs is None: - client = _flyte_engine.get_client() - execution_data = client.get_task_execution_data(self.id) - - # Inputs are returned inline unless they are too big, in which case a url blob pointing to them is returned. - if bool(execution_data.full_outputs.literals): - output_map = execution_data.full_outputs - - elif execution_data.outputs.bytes > 0: - with _common_utils.AutoDeletingTempDir() as t: - tmp_name = _os.path.join(t.name, "outputs.pb") - _data_proxy.Data.get_data(execution_data.outputs.url, tmp_name) - output_map = _literal_models.LiteralMap.from_flyte_idl( - _common_utils.load_proto_from_file(_literals_pb2.LiteralMap, tmp_name) - ) - else: - output_map = _literal_models.LiteralMap({}) - self._outputs = _type_helpers.unpack_literal_map_to_sdk_python_std(output_map) - return self._outputs - - @property - def error(self): - """ - If execution is in progress, raise an exception. Otherwise, return None if no error was present upon - reaching completion. - :rtype: flytekit.models.core.execution.ExecutionError or None - """ - if not self.is_complete: - raise _user_exceptions.FlyteAssertion( - "Please what until the task execution has completed before requesting error information." - ) - return self.closure.error - - def get_child_executions(self, filters=None): - """ - :param list[flytekit.models.filters.Filter] filters: - :rtype: dict[Text, flytekit.common.nodes.SdkNodeExecution] - """ - from flytekit.common import nodes as _nodes - - if not self.is_parent: - raise _user_exceptions.FlyteAssertion("Only task executions marked with 'is_parent' have child executions.") - client = _flyte_engine.get_client() - models = { - v.id.node_id: v - for v in _iterate_node_executions(client, task_execution_identifier=self.id, filters=filters) - } - - return {k: _nodes.SdkNodeExecution.promote_from_model(v) for k, v in _six.iteritems(models)} - - @classmethod - def promote_from_model(cls, base_model): - """ - :param _task_execution_model.TaskExecution base_model: - :rtype: SdkTaskExecution - """ - return cls( - closure=base_model.closure, - id=base_model.id, - input_uri=base_model.input_uri, - is_parent=base_model.is_parent, - ) - - def sync(self): - self._sync_closure() - - def _sync_closure(self): - """ - Syncs the closure of the underlying execution artifact with the state observed by the platform. - :rtype: None - """ - client = _flyte_engine.get_client() - self._closure = client.get_task_execution(self.id).closure diff --git a/flytekit/common/tasks/generic_spark_task.py b/flytekit/common/tasks/generic_spark_task.py deleted file mode 100644 index a83a7afbe4..0000000000 --- a/flytekit/common/tasks/generic_spark_task.py +++ /dev/null @@ -1,147 +0,0 @@ -import sys as _sys - -import six as _six -from google.protobuf.json_format import MessageToDict as _MessageToDict - -from flytekit import __version__ -from flytekit.common import interface as _interface -from flytekit.common.exceptions import scopes as _exception_scopes -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.tasks import task as _base_tasks -from flytekit.common.types import helpers as _helpers -from flytekit.common.types import primitives as _primitives -from flytekit.configuration import internal as _internal_config -from flytekit.models import literals as _literal_models -from flytekit.models import task as _task_models - -input_types_supported = { - _primitives.Integer, - _primitives.Boolean, - _primitives.Float, - _primitives.String, - _primitives.Datetime, - _primitives.Timedelta, -} - - -class SdkGenericSparkTask(_base_tasks.SdkTask): - """ - This class includes the additional logic for building a task that executes as a Spark Job. - - """ - - def __init__( - self, - task_type, - discovery_version, - retries, - interruptible, - task_inputs, - deprecated, - discoverable, - timeout, - spark_type, - main_class, - main_application_file, - spark_conf, - hadoop_conf, - environment, - ): - """ - :param Text task_type: string describing the task type - :param Text discovery_version: string describing the version for task discovery purposes - :param int retries: Number of retries to attempt - :param bool interruptible: Whether or not task is interruptible - :param Text deprecated: - :param bool discoverable: - :param datetime.timedelta timeout: - :param Text spark_type: Type of Spark Job: Scala/Java - :param Text main_class: Main class to execute for Scala/Java jobs - :param Text main_application_file: Main application file - :param dict[Text,Text] spark_conf: - :param dict[Text,Text] hadoop_conf: - :param dict[Text,Text] environment: [optional] environment variables to set when executing this task. - """ - - spark_job = _task_models.SparkJob( - spark_conf=spark_conf, - hadoop_conf=hadoop_conf, - spark_type=spark_type, - application_file=main_application_file, - main_class=main_class, - executor_path=_sys.executable, - ).to_flyte_idl() - - super(SdkGenericSparkTask, self).__init__( - task_type, - _task_models.TaskMetadata( - discoverable, - _task_models.RuntimeMetadata( - _task_models.RuntimeMetadata.RuntimeType.FLYTE_SDK, - __version__, - "spark", - ), - timeout, - _literal_models.RetryStrategy(retries), - interruptible, - discovery_version, - deprecated, - ), - _interface.TypedInterface({}, {}), - _MessageToDict(spark_job), - ) - - # Add Inputs - if task_inputs is not None: - task_inputs(self) - - # Container after the Inputs have been updated. - self._container = self._get_container_definition(environment=environment) - - def _validate_inputs(self, inputs): - """ - :param dict[Text, flytekit.models.interface.Variable] inputs: Input variables to validate - :raises: flytekit.common.exceptions.user.FlyteValidationException - """ - for k, v in _six.iteritems(inputs): - sdk_type = _helpers.get_sdk_type_from_literal_type(v.type) - if sdk_type not in input_types_supported: - raise _user_exceptions.FlyteValidationException( - "Input Type '{}' not supported. Only Primitives are supported for Scala/Java Spark.".format( - sdk_type - ) - ) - super(SdkGenericSparkTask, self)._validate_inputs(inputs) - - @_exception_scopes.system_entry_point - def add_inputs(self, inputs): - """ - Adds the inputs to this task. This can be called multiple times, but it will fail if an input with a given - name is added more than once, a name collides with an output, or if the name doesn't exist as an arg name in - the wrapped function. - :param dict[Text, flytekit.models.interface.Variable] inputs: names and variables - """ - self._validate_inputs(inputs) - self.interface.inputs.update(inputs) - - def _get_container_definition( - self, - environment=None, - ): - """ - :rtype: Container - """ - - args = [] - for k, v in _six.iteritems(self.interface.inputs): - args.append("--{}".format(k)) - args.append("{{{{.Inputs.{}}}}}".format(k)) - - return _task_models.Container( - image=_internal_config.IMAGE.get(), - command=[], - args=args, - resources=_task_models.Resources([], []), - env=environment, - config={}, - ) diff --git a/flytekit/common/tasks/hive_task.py b/flytekit/common/tasks/hive_task.py deleted file mode 100644 index 77ae3359a0..0000000000 --- a/flytekit/common/tasks/hive_task.py +++ /dev/null @@ -1,299 +0,0 @@ -import uuid as _uuid - -import six as _six -from google.protobuf.json_format import MessageToDict as _MessageToDict - -from flytekit.common import constants as _constants -from flytekit.common import interface as _interface -from flytekit.common import nodes as _nodes -from flytekit.common.exceptions import scopes as _exception_scopes -from flytekit.common.exceptions.user import FlyteTypeException as _FlyteTypeException -from flytekit.common.exceptions.user import FlyteValueException as _FlyteValueException -from flytekit.common.tasks import output as _task_output -from flytekit.common.tasks import sdk_runnable as _sdk_runnable -from flytekit.common.tasks import task as _base_task -from flytekit.common.types import helpers as _type_helpers -from flytekit.models import dynamic_job as _dynamic_job -from flytekit.models import interface as _interface_model -from flytekit.models import literals as _literal_models -from flytekit.models import qubole as _qubole -from flytekit.models.core import workflow as _workflow_model - -ALLOWED_TAGS_COUNT = int(6) -MAX_TAG_LENGTH = int(20) - - -class SdkHiveTask(_sdk_runnable.SdkRunnableTask): - """ - This class includes the additional logic for building a task that executes as a batch hive task. - """ - - def __init__( - self, - task_function, - task_type, - discovery_version, - retries, - interruptible, - deprecated, - storage_request, - cpu_request, - gpu_request, - memory_request, - storage_limit, - cpu_limit, - gpu_limit, - memory_limit, - discoverable, - timeout, - cluster_label, - tags, - environment, - cache_serializable, - ): - """ - :param task_function: Function container user code. This will be executed via the SDK's engine. - :param Text task_type: string describing the task type - :param Text discovery_version: string describing the version for task discovery purposes - :param int retries: Number of retries to attempt - :param Text deprecated: - :param Text storage_request: - :param Text cpu_request: - :param Text gpu_request: - :param Text memory_request: - :param Text storage_limit: - :param Text cpu_limit: - :param Text gpu_limit: - :param Text memory_limit: - :param bool discoverable: - :param datetime.timedelta timeout: - :param Text cluster_label: - :param list[Text] tags: - :param dict[Text, Text] environment: - :param bool cache_serializable: - """ - self._task_function = task_function - super(SdkHiveTask, self).__init__( - task_function, - task_type, - discovery_version, - retries, - interruptible, - deprecated, - storage_request, - cpu_request, - gpu_request, - memory_request, - storage_limit, - cpu_limit, - gpu_limit, - memory_limit, - discoverable, - timeout, - environment, - cache_serializable, - {}, - ) - self._validate_task_parameters(cluster_label, tags) - self._cluster_label = cluster_label - self._tags = tags - - def _generate_plugin_objects(self, context, inputs_dict): - """ - Runs user code and and produces hive queries - :param flytekit.engines.common.EngineContext context: - :param dict[Text, T] inputs: - :rtype: list[_qubole.QuboleHiveJob] - """ - queries_from_task = super(SdkHiveTask, self)._execute_user_code(context, inputs_dict) or [] - if not isinstance(queries_from_task, list): - queries_from_task = [queries_from_task] - - self._validate_queries(queries_from_task) - plugin_objects = [] - - for q in queries_from_task: - hive_query = _qubole.HiveQuery( - query=q, - timeout_sec=self.metadata.timeout.seconds, - retry_count=self.metadata.retries.retries, - ) - - # TODO: Remove this after all users of older SDK versions that did the single node, multi-query pattern are - # deprecated. This is only here for backwards compatibility - in addition to writing the query to the - # query field, we also construct a QueryCollection with only one query. This will ensure that the - # older plugin will continue to work. - query_collection = _qubole.HiveQueryCollection([hive_query]) - - plugin_objects.append( - _qubole.QuboleHiveJob( - hive_query, - self._cluster_label, - self._tags, - query_collection=query_collection, - ) - ) - - return plugin_objects - - @staticmethod - def _validate_task_parameters(cluster_label, tags): - if not (cluster_label is None or isinstance(cluster_label, (str, _six.text_type))): - raise _FlyteTypeException( - type(cluster_label), - {str, _six.text_type}, - additional_msg="cluster_label for a hive task must be in text format", - received_value=cluster_label, - ) - if tags is not None: - if not (isinstance(tags, list) and all(isinstance(tag, (str, _six.text_type)) for tag in tags)): - raise _FlyteTypeException( - type(tags), - [], - additional_msg="tags for a hive task must be in 'list of text' format", - received_value=tags, - ) - if len(tags) > ALLOWED_TAGS_COUNT: - raise _FlyteValueException( - len(tags), - "number of tags must be less than {}".format(ALLOWED_TAGS_COUNT), - ) - if not all(len(tag) for tag in tags): - raise _FlyteValueException( - tags, - "length of a tag must be less than {} chars".format(MAX_TAG_LENGTH), - ) - - @staticmethod - def _validate_queries(queries_from_task): - for query_from_task in queries_from_task or []: - if not isinstance(query_from_task, (str, _six.text_type)): - raise _FlyteTypeException( - type(query_from_task), - {str, _six.text_type}, - additional_msg="All queries returned from a Hive task must be in text format.", - received_value=query_from_task, - ) - - def _produce_dynamic_job_spec(self, context, inputs): - """ - Runs user code and and produces future task nodes to run sub-tasks. - :param context: - :param flytekit.models.literals.LiteralMap literal_map inputs: - :rtype: flytekit.models.dynamic_job.DynamicJobSpec - """ - inputs_dict = _type_helpers.unpack_literal_map_to_sdk_python_std( - inputs, - {k: _type_helpers.get_sdk_type_from_literal_type(v.type) for k, v in _six.iteritems(self.interface.inputs)}, - ) - outputs_dict = { - name: _task_output.OutputReference(_type_helpers.get_sdk_type_from_literal_type(variable.type)) - for name, variable in _six.iteritems(self.interface.outputs) - } - - # Add outputs to inputs - inputs_dict.update(outputs_dict) - - nodes = [] - tasks = [] - # One node per query - generated_queries = self._generate_plugin_objects(context, inputs_dict) - - # Create output bindings always - this has to happen after user code has run - output_bindings = [ - _literal_models.Binding( - var=name, - binding=_interface.BindingData.from_python_std(b.sdk_type.to_flyte_literal_type(), b.value), - ) - for name, b in _six.iteritems(outputs_dict) - ] - - i = 0 - for quboleHiveJob in generated_queries: - hive_job_node = _create_hive_job_node("HiveQuery_{}".format(i), quboleHiveJob.to_flyte_idl(), self.metadata) - nodes.append(hive_job_node) - tasks.append(hive_job_node.executable_sdk_object) - i += 1 - - dynamic_job_spec = _dynamic_job.DynamicJobSpec( - min_successes=len(nodes), - tasks=tasks, - nodes=nodes, - outputs=output_bindings, - subworkflows=[], - ) - - return dynamic_job_spec - - @_exception_scopes.system_entry_point - def execute(self, context, inputs): - """ - Executes hive batch task's user code and produces futures file as well as all sub-task inputs.pb files. - - :param flytekit.engines.common.EngineContext context: - :param flytekit.models.literals.LiteralMap inputs: - :rtype: dict[Text, flytekit.models.common.FlyteIdlEntity] - :returns: This function must return a dictionary mapping 'filenames' to Flyte Interface Entities. These - entities will be used by the engine to pass data from node to node, populate metadata, etc. etc.. Each - engine will have different behavior. For instance, the Flyte engine will upload the entities to a remote - working directory (with the names provided), which will in turn allow Flyte Propeller to push along the - workflow. Where as local engine will merely feed the outputs directly into the next node. - """ - spec = self._produce_dynamic_job_spec(context, inputs) - generated_files = {} - - # If no queries were produced, then the spec should not have any nodes, in which case we just produce an - # outputs file like any other single-step tasks. - if len(spec.nodes) == 0: - return { - _constants.OUTPUT_FILE_NAME: _literal_models.LiteralMap( - literals={binding.var: binding.binding.to_literal_model() for binding in spec.outputs} - ) - } - else: - generated_files.update({_constants.FUTURES_FILE_NAME: spec}) - - return generated_files - - -def _create_hive_job_node(name, hive_job, metadata): - """ - :param Text name: - :param _qubole.QuboleHiveJob hive_job: Hive job spec - :param flytekit.models.task.TaskMetadata metadata: This contains information needed at runtime to determine - behavior such as whether or not outputs are discoverable, timeouts, and retries. - :rtype: _nodes.SdkNode: - """ - return _nodes.SdkNode( - id=_six.text_type(_uuid.uuid4()), - upstream_nodes=[], - bindings=[], - metadata=_workflow_model.NodeMetadata(name, metadata.timeout, _literal_models.RetryStrategy(0)), - sdk_task=SdkHiveJob(hive_job, metadata), - ) - - -class SdkHiveJob(_base_task.SdkTask): - """ - This class encapsulates the hive-job that is submitted to the Qubole Operator. - - """ - - def __init__( - self, - hive_job, - metadata, - ): - """ - :param _qubole.QuboleHiveJob hive_job: Hive job spec - :param TaskMetadata metadata: This contains information needed at runtime to determine behavior such as - whether or not outputs are discoverable, timeouts, and retries. - """ - super(SdkHiveJob, self).__init__( - _constants.SdkTaskType.HIVE_JOB, - metadata, - # Individual hive tasks never take anything, or return anything. They just run a query that's already - # got the location set. - _interface_model.TypedInterface({}, {}), - _MessageToDict(hive_job), - ) diff --git a/flytekit/common/tasks/output.py b/flytekit/common/tasks/output.py deleted file mode 100644 index 8ca1b307cb..0000000000 --- a/flytekit/common/tasks/output.py +++ /dev/null @@ -1,46 +0,0 @@ -from flytekit.common.exceptions import scopes as _exception_scopes -from flytekit.common.types import base_sdk_types as _base_sdk_types - - -class OutputReference(object): - def __init__(self, sdk_type): - """ - :param flytekit.common.types.base_sdk_types.FlyteSdkType sdk_type: - """ - self._raw_value = None - self._sdk_type = sdk_type - self._sdk_value = _base_sdk_types.Void() - - @property - def value(self): - """ - :rtype: T - """ - return self._raw_value - - @property - def sdk_value(self): - """ - :rtype: flytekit.common.types.base_sdk_types.FlyteSdkValue - """ - return self._sdk_value - - @property - def sdk_type(self): - """ - :rtype: flytekit.common.types.base_sdk_types.FlyteSdkType - """ - return self._sdk_type - - @_exception_scopes.system_entry_point - def set(self, value): - """ - This should be called to set the value for output. The SDK will apply the appropriate type and value checking. - It will raise an exception if necessary. - :param T value: - :raises: flytekit.common.exceptions.user.FlyteValueException - """ - - sdk_value = self._sdk_type.from_python_std(value) - self._raw_value = value - self._sdk_value = sdk_value diff --git a/flytekit/common/tasks/presto_task.py b/flytekit/common/tasks/presto_task.py deleted file mode 100644 index c8f7d300a5..0000000000 --- a/flytekit/common/tasks/presto_task.py +++ /dev/null @@ -1,180 +0,0 @@ -import datetime as _datetime - -import six as _six -from google.protobuf.json_format import MessageToDict as _MessageToDict - -from flytekit import __version__ -from flytekit.common import constants as _constants -from flytekit.common import interface as _interface -from flytekit.common.exceptions import scopes as _exception_scopes -from flytekit.common.tasks import task as _base_task -from flytekit.common.types import helpers as _type_helpers -from flytekit.models import interface as _interface_model -from flytekit.models import literals as _literals -from flytekit.models import presto as _presto_models -from flytekit.models import task as _task_model -from flytekit.models import types as _types - - -class SdkPrestoTask(_base_task.SdkTask): - """ - This class includes the logic for building a task that executes as a Presto task. - """ - - def __init__( - self, - statement, - output_schema, - routing_group=None, - catalog=None, - schema=None, - task_inputs=None, - interruptible=False, - discoverable=False, - discovery_version=None, - retries=1, - timeout=None, - deprecated=None, - cache_serializable=False, - ): - """ - :param Text statement: Presto query specification - :param flytekit.common.types.schema.Schema output_schema: Schema that represents that data queried from Presto - :param Text routing_group: The routing group that a Presto query should be sent to for the given environment - :param Text catalog: The catalog to set for the given Presto query - :param Text schema: The schema to set for the given Presto query - :param dict[Text,flytekit.common.types.base_sdk_types.FlyteSdkType] task_inputs: Optional inputs to the Presto task - :param bool discoverable: - :param Text discovery_version: String describing the version for task discovery purposes - :param int retries: Number of retries to attempt - :param datetime.timedelta timeout: - :param Text deprecated: This string can be used to mark the task as deprecated. Consumers of the task will - receive deprecation warnings. - :param bool cache_serializable: - """ - - # Set as class fields which are used down below to configure implicit - # parameters - self._routing_group = routing_group or "" - self._catalog = catalog or "" - self._schema = schema or "" - - metadata = _task_model.TaskMetadata( - discoverable, - # This needs to have the proper version reflected in it - _task_model.RuntimeMetadata(_task_model.RuntimeMetadata.RuntimeType.FLYTE_SDK, __version__, "python"), - timeout or _datetime.timedelta(seconds=0), - _literals.RetryStrategy(retries), - interruptible, - discovery_version, - deprecated, - cache_serializable, - ) - - presto_query = _presto_models.PrestoQuery( - routing_group=routing_group or "", - catalog=catalog or "", - schema=schema or "", - statement=statement, - ) - - # Here we set the routing_group, catalog, and schema as implicit - # parameters for caching purposes - i = _interface.TypedInterface( - { - "__implicit_routing_group": _interface_model.Variable( - type=_types.LiteralType(simple=_types.SimpleType.STRING), - description="The routing group set as an implicit input", - ), - "__implicit_catalog": _interface_model.Variable( - type=_types.LiteralType(simple=_types.SimpleType.STRING), - description="The catalog set as an implicit input", - ), - "__implicit_schema": _interface_model.Variable( - type=_types.LiteralType(simple=_types.SimpleType.STRING), - description="The schema set as an implicit input", - ), - }, - { - # Set the schema for the Presto query as an output - "results": _interface_model.Variable( - type=_types.LiteralType(schema=output_schema.schema_type), - description="The schema for the Presto query", - ) - }, - ) - - super(SdkPrestoTask, self).__init__( - _constants.SdkTaskType.PRESTO_TASK, - metadata, - i, - _MessageToDict(presto_query.to_flyte_idl()), - ) - - # Set user provided inputs - task_inputs(self) - - def _add_implicit_inputs(self, inputs): - """ - :param dict[Text,Any] inputs: - :param inputs: - :return: - """ - inputs["__implicit_routing_group"] = self.routing_group - inputs["__implicit_catalog"] = self.catalog - inputs["__implicit_schema"] = self.schema - return inputs - - # Override method in order to set the implicit inputs - def __call__(self, *args, **kwargs): - kwargs = self._add_implicit_inputs(kwargs) - - return super(SdkPrestoTask, self).__call__(*args, **kwargs) - - # Override method in order to set the implicit inputs - def _python_std_input_map_to_literal_map(self, inputs): - """ - :param dict[Text,Any] inputs: A dictionary of Python standard inputs that will be type-checked and compiled - to a LiteralMap - :rtype: flytekit.models.literals.LiteralMap - """ - inputs = self._add_implicit_inputs(inputs) - return _type_helpers.pack_python_std_map_to_literal_map( - inputs, - {k: _type_helpers.get_sdk_type_from_literal_type(v.type) for k, v in _six.iteritems(self.interface.inputs)}, - ) - - @_exception_scopes.system_entry_point - def add_inputs(self, inputs): - """ - Adds the inputs to this task. This can be called multiple times, but it will fail if an input with a given - name is added more than once, a name collides with an output, or if the name doesn't exist as an arg name in - the wrapped function. - :param dict[Text, flytekit.models.interface.Variable] inputs: names and variables - """ - self._validate_inputs(inputs) - self.interface.inputs.update(inputs) - - @property - def routing_group(self): - """ - The routing group that a Presto query should be sent to for the given environment - :rtype: Text - """ - return self._routing_group - - @property - def catalog(self): - """ - The catalog to set for the given Presto query - :rtype: Text - """ - return self._catalog - - @property - def schema(self): - """ - The schema to set for the given Presto query - :rtype: Text - """ - return self._schema diff --git a/flytekit/common/tasks/raw_container.py b/flytekit/common/tasks/raw_container.py deleted file mode 100644 index 5168e744f3..0000000000 --- a/flytekit/common/tasks/raw_container.py +++ /dev/null @@ -1,237 +0,0 @@ -import datetime as _datetime -from typing import Dict, List - -import flytekit -from flytekit.common import constants as _constants -from flytekit.common import interface as _interface -from flytekit.common.exceptions import scopes as _exception_scopes -from flytekit.common.tasks import task as _base_task -from flytekit.common.types.base_sdk_types import FlyteSdkType -from flytekit.configuration import resources as _resource_config -from flytekit.models import literals as _literals -from flytekit.models import task as _task_models -from flytekit.models.interface import Variable - - -def types_to_variable(t: Dict[str, FlyteSdkType]) -> Dict[str, Variable]: - var = {} - if t: - for k, v in t.items(): - var[k] = Variable(v.to_flyte_literal_type(), "") - return var - - -def _get_container_definition( - image: str, - command: List[str], - args: List[str], - data_loading_config: _task_models.DataLoadingConfig, - storage_request: str = None, - ephemeral_storage_request: str = None, - cpu_request: str = None, - gpu_request: str = None, - memory_request: str = None, - storage_limit: str = None, - ephemeral_storage_limit: str = None, - cpu_limit: str = None, - gpu_limit: str = None, - memory_limit: str = None, - environment: Dict[str, str] = None, -) -> _task_models.Container: - storage_limit = storage_limit or _resource_config.DEFAULT_STORAGE_LIMIT.get() - storage_request = storage_request or _resource_config.DEFAULT_STORAGE_REQUEST.get() - ephemeral_storage_limit = ephemeral_storage_limit or _resource_config.DEFAULT_EPHEMERAL_STORAGE_LIMIT.get() - ephemeral_storage_request = ephemeral_storage_request or _resource_config.DEFAULT_EPHEMERAL_STORAGE_REQUEST.get() - cpu_limit = cpu_limit or _resource_config.DEFAULT_CPU_LIMIT.get() - cpu_request = cpu_request or _resource_config.DEFAULT_CPU_REQUEST.get() - gpu_limit = gpu_limit or _resource_config.DEFAULT_GPU_LIMIT.get() - gpu_request = gpu_request or _resource_config.DEFAULT_GPU_REQUEST.get() - memory_limit = memory_limit or _resource_config.DEFAULT_MEMORY_LIMIT.get() - memory_request = memory_request or _resource_config.DEFAULT_MEMORY_REQUEST.get() - - requests = [] - if storage_request: - requests.append( - _task_models.Resources.ResourceEntry(_task_models.Resources.ResourceName.STORAGE, storage_request) - ) - if ephemeral_storage_request: - requests.append( - _task_models.Resources.ResourceEntry( - _task_models.Resources.ResourceName.EPHEMERAL_STORAGE, ephemeral_storage_request - ) - ) - if cpu_request: - requests.append(_task_models.Resources.ResourceEntry(_task_models.Resources.ResourceName.CPU, cpu_request)) - if gpu_request: - requests.append(_task_models.Resources.ResourceEntry(_task_models.Resources.ResourceName.GPU, gpu_request)) - if memory_request: - requests.append( - _task_models.Resources.ResourceEntry(_task_models.Resources.ResourceName.MEMORY, memory_request) - ) - - limits = [] - if storage_limit: - limits.append(_task_models.Resources.ResourceEntry(_task_models.Resources.ResourceName.STORAGE, storage_limit)) - if ephemeral_storage_limit: - limits.append( - _task_models.Resources.ResourceEntry( - _task_models.Resources.ResourceName.EPHEMERAL_STORAGE, ephemeral_storage_limit - ) - ) - if cpu_limit: - limits.append(_task_models.Resources.ResourceEntry(_task_models.Resources.ResourceName.CPU, cpu_limit)) - if gpu_limit: - limits.append(_task_models.Resources.ResourceEntry(_task_models.Resources.ResourceName.GPU, gpu_limit)) - if memory_limit: - limits.append(_task_models.Resources.ResourceEntry(_task_models.Resources.ResourceName.MEMORY, memory_limit)) - - if environment is None: - environment = {} - - return _task_models.Container( - image=image, - command=command, - args=args, - resources=_task_models.Resources(limits=limits, requests=requests), - env=environment, - config={}, - data_loading_config=data_loading_config, - ) - - -class SdkRawContainerTask(_base_task.SdkTask): - """ - Use this task when you want to run an arbitrary container as a task (e.g. external tools, binaries compiled - separately as a container completely separate from the container where your Flyte workflow is defined. - """ - - METADATA_FORMAT_JSON = _task_models.DataLoadingConfig.LITERALMAP_FORMAT_JSON - METADATA_FORMAT_YAML = _task_models.DataLoadingConfig.LITERALMAP_FORMAT_YAML - METADATA_FORMAT_PROTO = _task_models.DataLoadingConfig.LITERALMAP_FORMAT_PROTO - - def __init__( - self, - inputs: Dict[str, FlyteSdkType], - image: str, - outputs: Dict[str, FlyteSdkType] = None, - input_data_dir: str = None, - output_data_dir: str = None, - metadata_format: int = METADATA_FORMAT_JSON, - io_strategy: _task_models.IOStrategy = None, - command: List[str] = None, - args: List[str] = None, - storage_request: str = None, - cpu_request: str = None, - gpu_request: str = None, - memory_request: str = None, - storage_limit: str = None, - cpu_limit: str = None, - gpu_limit: str = None, - memory_limit: str = None, - environment: Dict[str, str] = None, - interruptible: bool = False, - discoverable: bool = False, - discovery_version: str = None, - retries: int = 1, - timeout: _datetime.timedelta = None, - cache_serializable: bool = False, - ): - """ - :param inputs: - :param outputs: - :param image: - :param command: - :param args: - :param storage_request: - :param cpu_request: - :param gpu_request: - :param memory_request: - :param storage_limit: - :param cpu_limit: - :param gpu_limit: - :param memory_limit: - :param environment: - :param interruptible: - :param discoverable: - :param discovery_version: - :param retries: - :param timeout: - :param cache_serializable: - :param input_data_dir: This is the directory where data will be downloaded to - :param output_data_dir: This is the directory where data will be uploaded from - :param metadata_format: Format in which the metadata will be available for the script - """ - - # Set as class fields which are used down below to configure implicit - # parameters - self._data_loading_config = _task_models.DataLoadingConfig( - input_path=input_data_dir, - output_path=output_data_dir, - format=metadata_format, - enabled=True, - io_strategy=io_strategy, - ) - - metadata = _task_models.TaskMetadata( - discoverable, - # This needs to have the proper version reflected in it - _task_models.RuntimeMetadata( - _task_models.RuntimeMetadata.RuntimeType.FLYTE_SDK, - flytekit.__version__, - "python", - ), - timeout or _datetime.timedelta(seconds=0), - _literals.RetryStrategy(retries), - interruptible, - discovery_version, - None, - cache_serializable, - ) - - # The interface is defined using the inputs and outputs - i = _interface.TypedInterface(inputs=types_to_variable(inputs), outputs=types_to_variable(outputs)) - - # This sets the base SDKTask with container etc - super(SdkRawContainerTask, self).__init__( - _constants.SdkTaskType.RAW_CONTAINER_TASK, - metadata, - i, - None, - container=_get_container_definition( - image=image, - args=args, - command=command, - data_loading_config=self._data_loading_config, - storage_request=storage_request, - cpu_request=cpu_request, - gpu_request=gpu_request, - memory_request=memory_request, - storage_limit=storage_limit, - cpu_limit=cpu_limit, - gpu_limit=gpu_limit, - memory_limit=memory_limit, - environment=environment, - ), - ) - - @_exception_scopes.system_entry_point - def add_inputs(self, inputs: Dict[str, Variable]): - """ - Adds the inputs to this task. This can be called multiple times, but it will fail if an input with a given - name is added more than once, a name collides with an output, or if the name doesn't exist as an arg name in - the wrapped function. - :param dict[Text, flytekit.models.interface.Variable] inputs: names and variables - """ - self._validate_inputs(inputs) - self.interface.inputs.update(inputs) - - @_exception_scopes.system_entry_point - def add_outputs(self, outputs: Dict[str, Variable]): - """ - Adds the inputs to this task. This can be called multiple times, but it will fail if an input with a given - name is added more than once, a name collides with an output, or if the name doesn't exist as an arg name in - the wrapped function. - :param dict[Text, flytekit.models.interface.Variable] outputs: names and variables - """ - self._validate_outputs(outputs) - self.interface.outputs.update(outputs) diff --git a/flytekit/common/tasks/sdk_dynamic.py b/flytekit/common/tasks/sdk_dynamic.py deleted file mode 100644 index b3bc923601..0000000000 --- a/flytekit/common/tasks/sdk_dynamic.py +++ /dev/null @@ -1,372 +0,0 @@ -import itertools as _itertools -import math -import os as _os - -import six as _six - -from flytekit.common import constants as _constants -from flytekit.common import interface as _interface -from flytekit.common import launch_plan as _launch_plan -from flytekit.common import sdk_bases as _sdk_bases -from flytekit.common import workflow as _workflow -from flytekit.common.core import identifier as _identifier -from flytekit.common.exceptions import scopes as _exception_scopes -from flytekit.common.mixins import registerable as _registerable -from flytekit.common.tasks import output as _task_output -from flytekit.common.tasks import sdk_runnable as _sdk_runnable -from flytekit.common.tasks import task as _task -from flytekit.common.types import helpers as _type_helpers -from flytekit.common.utils import _dnsify -from flytekit.configuration import internal as _internal_config -from flytekit.models import array_job as _array_job -from flytekit.models import dynamic_job as _dynamic_job -from flytekit.models import literals as _literal_models - - -class PromiseOutputReference(_task_output.OutputReference): - @property - def raw_value(self): - """ - :rtype: T - """ - return self._raw_value - - @_exception_scopes.system_entry_point - def set(self, value): - """ - This should be called to set the value for output. The SDK will apply the appropriate type and value checking. - It will raise an exception if necessary. - :param T value: - :raises: flytekit.common.exceptions.user.FlyteValueException - """ - - self._raw_value = value - - -def _append_node(generated_files, node, nodes, sub_task_node): - nodes.append(node) - for k, node_output in _six.iteritems(sub_task_node.outputs): - if not node_output.sdk_node.id: - node_output.sdk_node.assign_id_and_return(node.id) - - # Upload inputs to working directory under /array_job.input_ref/inputs.pb - input_path = _os.path.join(node.id, _constants.INPUT_FILE_NAME) - generated_files[input_path] = _literal_models.LiteralMap( - literals={binding.var: binding.binding.to_literal_model() for binding in sub_task_node.inputs} - ) - - -class SdkDynamicTaskMixin(object): - - """ - This mixin implements logic for building a task that executes - parent-child tasks in Python code. - - """ - - def __init__(self, allowed_failure_ratio, max_concurrency): - """ - :param float allowed_failure_ratio: - :param int max_concurrency: - """ - - # These will only appear in the generated futures - self._allowed_failure_ratio = allowed_failure_ratio - self._max_concurrency = max_concurrency - - def _create_array_job(self, inputs_prefix): - """ - Creates an array job for the passed sdk_task. - :param str inputs_prefix: - :rtype: _array_job.ArrayJob - """ - return _array_job.ArrayJob( - parallelism=self._max_concurrency if self._max_concurrency else 0, - size=1, - min_successes=1, - ) - - @staticmethod - def _can_run_as_array(task_type): - """ - Checks if a task can be grouped to run as an array job. - :param Text task_type: - :rtype: bool - """ - return task_type == _constants.SdkTaskType.PYTHON_TASK - - @staticmethod - def _add_upstream_entities(executable_sdk_object, sub_workflows, tasks): - upstream_entities = [] - if isinstance(executable_sdk_object, _workflow.SdkWorkflow): - upstream_entities = [n.executable_sdk_object for n in executable_sdk_object.nodes] - - for upstream_entity in upstream_entities: - # If the upstream entity is either a Workflow or a Task, yield them in the - # dynamic job spec. Otherwise (e.g. a LaunchPlan), we will assume it already - # is registered (can't be dynamically created). This will cause a runtime error - # if it's not already registered with the control plane. - if isinstance(upstream_entity, _workflow.SdkWorkflow): - sub_workflows.add(upstream_entity) - # Recursively discover all statically defined dependencies - SdkDynamicTask._add_upstream_entities(upstream_entity, sub_workflows, tasks) - elif isinstance(upstream_entity, _task.SdkTask): - tasks.add(upstream_entity) - - def _produce_dynamic_job_spec(self, context, inputs): - """ - Runs user code and and produces future task nodes to run sub-tasks. - :param context: - :param flytekit.models.literals.LiteralMap literal_map inputs: - :rtype: (_dynamic_job.DynamicJobSpec, dict[Text, flytekit.models.common.FlyteIdlEntity]) - """ - inputs_dict = _type_helpers.unpack_literal_map_to_sdk_python_std( - inputs, - {k: _type_helpers.get_sdk_type_from_literal_type(v.type) for k, v in _six.iteritems(self.interface.inputs)}, - ) - outputs_dict = { - name: PromiseOutputReference(_type_helpers.get_sdk_type_from_literal_type(variable.type)) - for name, variable in _six.iteritems(self.interface.outputs) - } - - # Because users declare both inputs and outputs in their functions signatures, merge them together - # before calling user code - inputs_dict.update(outputs_dict) - yielded_sub_tasks = [sub_task for sub_task in self._execute_user_code(context, inputs_dict) or []] - - upstream_nodes = list() - output_bindings = [ - _literal_models.Binding( - var=name, - binding=_interface.BindingData.from_python_std( - b.sdk_type.to_flyte_literal_type(), - b.raw_value, - upstream_nodes=upstream_nodes, - ), - ) - for name, b in _six.iteritems(outputs_dict) - ] - upstream_nodes = set(upstream_nodes) - - generated_files = {} - # Keeping future-tasks in original order. We don't use upstream_nodes exclusively because the parent task can - # yield sub-tasks that it never uses to produce final outputs but they need to execute nevertheless. - array_job_index = {} - tasks = set() - nodes = [] - sub_workflows = set() - visited_nodes = set() - generated_ids = {} - effective_failure_ratio = self._allowed_failure_ratio or 0.0 - - # TODO: This function needs to be cleaned up. - # The reason we chain these two together is because we allow users to not have to explicitly "yield" the - # node. As long as the subtask/lp/subwf has an output that's referenced, it'll get picked up. - for sub_task_node in _itertools.chain(yielded_sub_tasks, upstream_nodes): - if sub_task_node in visited_nodes: - continue - visited_nodes.add(sub_task_node) - executable = sub_task_node.executable_sdk_object - - # If the executable object that we're dealing with is registerable (ie, SdkRunnableLaunchPlan, SdkWorkflow - # SdkTask, or SdkRunnableTask), then it should have the ability to give itself a name. After assigning - # itself the name, also make sure the id is properly set according to current config values. - if isinstance(executable, _registerable.TrackableEntity) and not executable.has_valid_name: - executable.auto_assign_name() - executable._id = _identifier.Identifier( - executable.resource_type, - _internal_config.TASK_PROJECT.get() or _internal_config.PROJECT.get(), - _internal_config.TASK_DOMAIN.get() or _internal_config.DOMAIN.get(), - executable.platform_valid_name, - _internal_config.TASK_VERSION.get() or _internal_config.VERSION.get(), - ) - - # Generate an id that's unique in the document (if the same task is used multiple times with - # different resources, executable_sdk_object.id will be the same but generated node_ids should not - # be. - safe_task_id = _six.text_type(sub_task_node.executable_sdk_object.id) - if safe_task_id in generated_ids: - new_count = generated_ids[safe_task_id] = generated_ids[safe_task_id] + 1 - else: - new_count = generated_ids[safe_task_id] = 0 - unique_node_id = _dnsify("{}-{}".format(safe_task_id, new_count)) - - # Handling case where the yielded node is launch plan - if isinstance(sub_task_node.executable_sdk_object, _launch_plan.SdkLaunchPlan): - node = sub_task_node.assign_id_and_return(unique_node_id) - _append_node(generated_files, node, nodes, sub_task_node) - # Handling case where the yielded node is launching a sub-workflow - elif isinstance(sub_task_node.executable_sdk_object, _workflow.SdkWorkflow): - node = sub_task_node.assign_id_and_return(unique_node_id) - _append_node(generated_files, node, nodes, sub_task_node) - # Add the workflow itself to the yielded sub-workflows - sub_workflows.add(sub_task_node.executable_sdk_object) - # Recursively discover statically defined upstream entities (tasks, wfs) - SdkDynamicTask._add_upstream_entities(sub_task_node.executable_sdk_object, sub_workflows, tasks) - # Handling tasks - else: - # If the task can run as an array job, group its instances together. Otherwise, keep each - # invocation as a separate node. - if SdkDynamicTask._can_run_as_array(sub_task_node.executable_sdk_object.type): - if sub_task_node.executable_sdk_object in array_job_index: - array_job, node = array_job_index[sub_task_node.executable_sdk_object] - array_job.size += 1 - array_job.min_successes = int(math.ceil((1 - effective_failure_ratio) * array_job.size)) - else: - array_job = self._create_array_job(inputs_prefix=unique_node_id) - node = sub_task_node.assign_id_and_return(unique_node_id) - array_job_index[sub_task_node.executable_sdk_object] = ( - array_job, - node, - ) - - node_index = _six.text_type(array_job.size - 1) - for k, node_output in _six.iteritems(sub_task_node.outputs): - if not node_output.sdk_node.id: - node_output.sdk_node.assign_id_and_return(node.id) - node_output.var = "[{}].{}".format(node_index, node_output.var) - - # Upload inputs to working directory under /array_job.input_ref//inputs.pb - input_path = _os.path.join(node.id, node_index, _constants.INPUT_FILE_NAME) - generated_files[input_path] = _literal_models.LiteralMap( - literals={binding.var: binding.binding.to_literal_model() for binding in sub_task_node.inputs} - ) - else: - node = sub_task_node.assign_id_and_return(unique_node_id) - tasks.add(sub_task_node.executable_sdk_object) - _append_node(generated_files, node, nodes, sub_task_node) - - # assign custom field to the ArrayJob properties computed. - for task, (array_job, _) in _six.iteritems(array_job_index): - # TODO: Reconstruct task template object instead of modifying an existing one? - tasks.add( - task.assign_custom_and_return(array_job.to_dict()).assign_type_and_return( - _constants.SdkTaskType.CONTAINER_ARRAY_TASK - ) - ) - - # min_successes is absolute, it's computed as the reverse of allowed_failure_ratio and multiplied by the - # total length of tasks to get an absolute count. - nodes.extend([array_job_node for (_, array_job_node) in array_job_index.values()]) - dynamic_job_spec = _dynamic_job.DynamicJobSpec( - min_successes=len(nodes), - tasks=list(tasks), - nodes=nodes, - outputs=output_bindings, - subworkflows=list(sub_workflows), - ) - - return dynamic_job_spec, generated_files - - @_exception_scopes.system_entry_point - def execute(self, context, inputs): - """ - Executes batch task's user code and produces futures file as well as all sub-task inputs.pb files. - - :param flytekit.engines.common.EngineContext context: - :param flytekit.models.literals.LiteralMap inputs: - :rtype: dict[Text, flytekit.models.common.FlyteIdlEntity] - :returns: This function must return a dictionary mapping 'filenames' to Flyte Interface Entities. These - entities will be used by the engine to pass data from node to node, populate metadata, etc. etc.. Each - engine will have different behavior. For instance, the Flyte engine will upload the entities to a remote - working directory (with the names provided), which will in turn allow Flyte Propeller to push along the - workflow. Where as local engine will merely feed the outputs directly into the next node. - """ - spec, generated_files = self._produce_dynamic_job_spec(context, inputs) - - # If no sub-tasks are requested to run, just produce an outputs file like any other single-step tasks. - if len(spec.nodes) == 0: - return { - _constants.OUTPUT_FILE_NAME: _literal_models.LiteralMap( - literals={binding.var: binding.binding.to_literal_model() for binding in spec.outputs} - ) - } - else: - generated_files.update({_constants.FUTURES_FILE_NAME: spec}) - - return generated_files - - -class SdkDynamicTask( - SdkDynamicTaskMixin, - _sdk_runnable.SdkRunnableTask, - metaclass=_sdk_bases.ExtendedSdkType, -): - - """ - This class includes the additional logic for building a task that executes - parent-child tasks in Python code. - - """ - - def __init__( - self, - task_function, - task_type, - discovery_version, - retries, - interruptible, - deprecated, - storage_request, - cpu_request, - gpu_request, - memory_request, - storage_limit, - cpu_limit, - gpu_limit, - memory_limit, - discoverable, - timeout, - allowed_failure_ratio, - max_concurrency, - environment, - cache_serializable, - custom, - ): - """ - :param task_function: Function container user code. This will be executed via the SDK's engine. - :param Text task_type: string describing the task type - :param Text discovery_version: string describing the version for task discovery purposes - :param int retries: Number of retries to attempt - :param bool interruptible: Whether or not task is interruptible - :param Text deprecated: - :param Text storage_request: - :param Text cpu_request: - :param Text gpu_request: - :param Text memory_request: - :param Text storage_limit: - :param Text cpu_limit: - :param Text gpu_limit: - :param Text memory_limit: - :param bool discoverable: - :param datetime.timedelta timeout: - :param float allowed_failure_ratio: - :param int max_concurrency: - :param dict[Text, Text] environment: - :param bool cache_serializable: - :param dict[Text, T] custom: - """ - _sdk_runnable.SdkRunnableTask.__init__( - self, - task_function, - task_type, - discovery_version, - retries, - interruptible, - deprecated, - storage_request, - cpu_request, - gpu_request, - memory_request, - storage_limit, - cpu_limit, - gpu_limit, - memory_limit, - discoverable, - timeout, - environment, - cache_serializable, - custom, - ) - - SdkDynamicTaskMixin.__init__(self, allowed_failure_ratio, max_concurrency) diff --git a/flytekit/common/tasks/sdk_runnable.py b/flytekit/common/tasks/sdk_runnable.py deleted file mode 100644 index 39788437dd..0000000000 --- a/flytekit/common/tasks/sdk_runnable.py +++ /dev/null @@ -1,750 +0,0 @@ -from __future__ import annotations - -import copy as _copy -import enum -import logging as _logging -import os -import pathlib -import typing -from dataclasses import dataclass -from datetime import datetime -from inspect import getfullargspec as _getargspec - -import six as _six - -from flytekit.common import constants as _constants -from flytekit.common import interface as _interface -from flytekit.common import sdk_bases as _sdk_bases -from flytekit.common import utils as _common_utils -from flytekit.common.core.identifier import WorkflowExecutionIdentifier -from flytekit.common.exceptions import scopes as _exception_scopes -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.tasks import output as _task_output -from flytekit.common.tasks import task as _base_task -from flytekit.common.types import helpers as _type_helpers -from flytekit.configuration import internal as _internal_config -from flytekit.configuration import resources as _resource_config -from flytekit.configuration import sdk as _sdk_config -from flytekit.configuration import secrets -from flytekit.engines import loader as _engine_loader -from flytekit.interfaces.stats import taggable -from flytekit.models import literals as _literal_models -from flytekit.models import task as _task_models - - -class SecretsManager(object): - """ - This provides a secrets resolution logic at runtime. - The resolution order is - - Try env var first. The env var should have the configuration.SECRETS_ENV_PREFIX. The env var will be all upper - cased - - If not then try the file where the name matches lower case - ``configuration.SECRETS_DEFAULT_DIR//configuration.SECRETS_FILE_PREFIX`` - - All configuration values can always be overridden by injecting an environment variable - """ - - def __init__(self): - self._base_dir = str(secrets.SECRETS_DEFAULT_DIR.get()).strip() - self._file_prefix = str(secrets.SECRETS_FILE_PREFIX.get()).strip() - self._env_prefix = str(secrets.SECRETS_ENV_PREFIX.get()).strip() - - def get(self, group: str, key: str) -> str: - """ - Retrieves a secret using the resolution order -> Env followed by file. If not found raises a ValueError - """ - self.check_group_key(group, key) - env_var = self.get_secrets_env_var(group, key) - fpath = self.get_secrets_file(group, key) - v = os.environ.get(env_var) - if v is not None: - return v - if os.path.exists(fpath): - with open(fpath, "r") as f: - return f.read().strip() - raise ValueError( - f"Unable to find secret for key {key} in group {group} " f"in Env Var:{env_var} and FilePath: {fpath}" - ) - - def get_secrets_env_var(self, group: str, key: str) -> str: - """ - Returns a string that matches the ENV Variable to look for the secrets - """ - self.check_group_key(group, key) - return f"{self._env_prefix}{group.upper()}_{key.upper()}" - - def get_secrets_file(self, group: str, key: str) -> str: - """ - Returns a path that matches the file to look for the secrets - """ - self.check_group_key(group, key) - return os.path.join(self._base_dir, group.lower(), f"{self._file_prefix}{key.lower()}") - - @staticmethod - def check_group_key(group: str, key: str): - if group is None or group == "": - raise ValueError("secrets group is a mandatory field.") - if key is None or key == "": - raise ValueError("secrets key is a mandatory field.") - - -# TODO: Clean up working dir name -class ExecutionParameters(object): - """ - This is a run-time user-centric context object that is accessible to every @task method. It can be accessed using - - .. code-block:: python - - flytekit.current_context() - - This object provides the following - * a statsd handler - * a logging handler - * the execution ID as an :py:class:`flytekit.models.core.identifier.WorkflowExecutionIdentifier` object - * a working directory for the user to write arbitrary files to - - Please do not confuse this object with the :py:class:`flytekit.FlyteContext` object. - """ - - @dataclass(init=False) - class Builder(object): - stats: taggable.TaggableStats - execution_date: datetime - logging: _logging - execution_id: str - attrs: typing.Dict[str, typing.Any] - working_dir: typing.Union[os.PathLike, _common_utils.AutoDeletingTempDir] - - def __init__(self, current: typing.Optional[ExecutionParameters] = None): - self.stats = current.stats if current else None - self.execution_date = current.execution_date if current else None - self.working_dir = current.working_directory if current else None - self.execution_id = current.execution_id if current else None - self.logging = current.logging if current else None - self.attrs = current._attrs if current else {} - - def add_attr(self, key: str, v: typing.Any) -> ExecutionParameters.Builder: - self.attrs[key] = v - return self - - def build(self) -> ExecutionParameters: - if not isinstance(self.working_dir, _common_utils.AutoDeletingTempDir): - pathlib.Path(self.working_dir).mkdir(parents=True, exist_ok=True) - return ExecutionParameters( - execution_date=self.execution_date, - stats=self.stats, - tmp_dir=self.working_dir, - execution_id=self.execution_id, - logging=self.logging, - **self.attrs, - ) - - @staticmethod - def new_builder(current: ExecutionParameters = None) -> Builder: - return ExecutionParameters.Builder(current=current) - - def builder(self) -> Builder: - return ExecutionParameters.Builder(current=self) - - def __init__(self, execution_date, tmp_dir, stats, execution_id, logging, **kwargs): - """ - Args: - execution_date: Date when the execution is running - tmp_dir: temporary directory for the execution - stats: handle to emit stats - execution_id: Identifier for the xecution - logging: handle to logging - """ - self._stats = stats - self._execution_date = execution_date - self._working_directory = tmp_dir - self._execution_id = execution_id - self._logging = logging - # AutoDeletingTempDir's should be used with a with block, which creates upon entry - self._attrs = kwargs - # It is safe to recreate the Secrets Manager - self._secrets_manager = SecretsManager() - - @property - def stats(self) -> taggable.TaggableStats: - """ - A handle to a special statsd object that provides usefully tagged stats. - TODO: Usage examples and better comments - """ - return self._stats - - @property - def logging(self) -> _logging: - """ - A handle to a useful logging object. - TODO: Usage examples - """ - return self._logging - - @property - def working_directory(self) -> _common_utils.AutoDeletingTempDir: - """ - A handle to a special working directory for easily producing temporary files. - - TODO: Usage examples - TODO: This does not always return a AutoDeletingTempDir - """ - return self._working_directory - - @property - def execution_date(self) -> datetime: - """ - This is a datetime representing the time at which a workflow was started. This is consistent across all tasks - executed in a workflow or sub-workflow. - - .. note:: - - Do NOT use this execution_date to drive any production logic. It might be useful as a tag for data to help - in debugging. - """ - return self._execution_date - - @property - def execution_id(self) -> str: - """ - This is the identifier of the workflow execution within the underlying engine. It will be consistent across all - task executions in a workflow or sub-workflow execution. - - .. note:: - - Do NOT use this execution_id to drive any production logic. This execution ID should only be used as a tag - on output data to link back to the workflow run that created it. - """ - return self._execution_id - - @property - def secrets(self) -> SecretsManager: - return self._secrets_manager - - def __getattr__(self, attr_name: str) -> typing.Any: - """ - This houses certain task specific context. For example in Spark, it houses the SparkSession, etc - """ - attr_name = attr_name.upper() - if self._attrs and attr_name in self._attrs: - return self._attrs[attr_name] - raise AssertionError(f"{attr_name} not available as a parameter in Flyte context - are you in right task-type?") - - def has_attr(self, attr_name: str) -> bool: - attr_name = attr_name.upper() - if self._attrs and attr_name in self._attrs: - return True - return False - - def get(self, key: str) -> typing.Any: - """ - Returns task specific context if present else raise an error. The returned context will match the key - """ - return self.__getattr__(attr_name=key) - - -class SdkRunnableContainer(_task_models.Container, metaclass=_sdk_bases.ExtendedSdkType): - """ - This is not necessarily a local-only Container object. So long as configuration is present, you can use this object - """ - - def __init__( - self, - command, - args, - resources, - env, - config, - ): - super(SdkRunnableContainer, self).__init__("", command, args, resources, env or {}, config) - - @property - def args(self): - """ - :rtype: list[Text] - """ - return _sdk_config.SDK_PYTHON_VENV.get() + self._args - - @property - def image(self): - """ - :rtype: Text - """ - return _internal_config.IMAGE.get() - - @property - def env(self): - """ - :rtype: dict[Text,Text] - """ - env = super(SdkRunnableContainer, self).env.copy() - env.update( - { - _internal_config.CONFIGURATION_PATH.env_var: _internal_config.CONFIGURATION_PATH.get(), - _internal_config.IMAGE.env_var: _internal_config.IMAGE.get(), - # TODO: Phase out the below. Propeller will set these and these are not SDK specific - _internal_config.PROJECT.env_var: _internal_config.PROJECT.get(), - _internal_config.DOMAIN.env_var: _internal_config.DOMAIN.get(), - _internal_config.NAME.env_var: _internal_config.NAME.get(), - _internal_config.VERSION.env_var: _internal_config.VERSION.get(), - } - ) - return env - - @classmethod - def get_resources( - cls, - storage_request=None, - cpu_request=None, - gpu_request=None, - memory_request=None, - storage_limit=None, - cpu_limit=None, - gpu_limit=None, - memory_limit=None, - ): - """ - :param Text storage_request: - :param Text cpu_request: - :param Text gpu_request: - :param Text memory_request: - :param Text storage_limit: - :param Text cpu_limit: - :param Text gpu_limit: - :param Text memory_limit: - """ - requests = [] - if storage_request: - requests.append( - _task_models.Resources.ResourceEntry(_task_models.Resources.ResourceName.STORAGE, storage_request) - ) - if cpu_request: - requests.append(_task_models.Resources.ResourceEntry(_task_models.Resources.ResourceName.CPU, cpu_request)) - if gpu_request: - requests.append(_task_models.Resources.ResourceEntry(_task_models.Resources.ResourceName.GPU, gpu_request)) - if memory_request: - requests.append( - _task_models.Resources.ResourceEntry(_task_models.Resources.ResourceName.MEMORY, memory_request) - ) - - limits = [] - if storage_limit: - limits.append( - _task_models.Resources.ResourceEntry(_task_models.Resources.ResourceName.STORAGE, storage_limit) - ) - if cpu_limit: - limits.append(_task_models.Resources.ResourceEntry(_task_models.Resources.ResourceName.CPU, cpu_limit)) - if gpu_limit: - limits.append(_task_models.Resources.ResourceEntry(_task_models.Resources.ResourceName.GPU, gpu_limit)) - if memory_limit: - limits.append( - _task_models.Resources.ResourceEntry(_task_models.Resources.ResourceName.MEMORY, memory_limit) - ) - - return _task_models.Resources(limits=limits, requests=requests) - - -class SdkRunnableTaskStyle(enum.Enum): - V0 = 0 - V1 = 1 - - -class SdkRunnableTask(_base_task.SdkTask, metaclass=_sdk_bases.ExtendedSdkType): - """ - This class includes the additional logic for building a task that executes in Python code. It has even more - validation checks to ensure proper behavior than it's superclasses. - - Since an SdkRunnableTask is assumed to run by hooking into Python code, we will provide additional shortcuts and - methods on this object. - """ - - def __init__( - self, - task_function, - task_type, - discovery_version, - retries, - interruptible, - deprecated, - storage_request, - cpu_request, - gpu_request, - memory_request, - storage_limit, - cpu_limit, - gpu_limit, - memory_limit, - discoverable, - timeout, - environment, - cache_serializable, - custom, - ): - """ - :param task_function: Function container user code. This will be executed via the SDK's engine. - :param Text task_type: string describing the task type - :param Text discovery_version: string describing the version for task discovery purposes - :param int retries: Number of retries to attempt - :param bool interruptible: Specify whether task is interruptible - :param Text deprecated: - :param Text storage_request: - :param Text cpu_request: - :param Text gpu_request: - :param Text memory_request: - :param Text storage_limit: - :param Text cpu_limit: - :param Text gpu_limit: - :param Text memory_limit: - :param bool discoverable: - :param datetime.timedelta timeout: - :param dict[Text, Text] environment: - :param bool cache_serializable: - :param dict[Text, T] custom: - """ - # Circular dependency - from flytekit import __version__ - - self._task_function = task_function - super(SdkRunnableTask, self).__init__( - task_type, - _task_models.TaskMetadata( - discoverable, - _task_models.RuntimeMetadata( - _task_models.RuntimeMetadata.RuntimeType.FLYTE_SDK, - __version__, - "python", - ), - timeout, - _literal_models.RetryStrategy(retries), - interruptible, - discovery_version, - deprecated, - cache_serializable, - ), - # TODO: If we end up using SdkRunnableTask for the new code, make sure this is set correctly. - _interface.TypedInterface({}, {}), - custom, - container=self._get_container_definition( - storage_request=storage_request, - cpu_request=cpu_request, - gpu_request=gpu_request, - memory_request=memory_request, - storage_limit=storage_limit, - cpu_limit=cpu_limit, - gpu_limit=gpu_limit, - memory_limit=memory_limit, - environment=environment, - ), - ) - self.id._name = "{}.{}".format(self.task_module, self.task_function_name) - self._has_fast_registered = False - - # TODO: Remove this in the future, I don't think we'll be using this. - self._task_style = SdkRunnableTaskStyle.V0 - - _banned_inputs = {} - _banned_outputs = {} - - @_exception_scopes.system_entry_point - def add_inputs(self, inputs): - """ - Adds the inputs to this task. This can be called multiple times, but it will fail if an input with a given - name is added more than once, a name collides with an output, or if the name doesn't exist as an arg name in - the wrapped function. - :param dict[Text, flytekit.models.interface.Variable] inputs: names and variables - """ - self._validate_inputs(inputs) - self.interface.inputs.update(inputs) - - @classmethod - def promote_from_model(cls, base_model): - # TODO: If the task exists in this container, we should be able to retrieve it. - raise _user_exceptions.FlyteAssertion("Cannot promote a base object to a runnable task.") - - @property - def task_style(self): - return self._task_style - - @property - def task_function(self): - return self._task_function - - @property - def task_function_name(self): - """ - :rtype: Text - """ - return self.task_function.__name__ - - @property - def task_module(self): - """ - :rtype: Text - """ - return self._task_function.__module__ - - def validate(self): - super(SdkRunnableTask, self).validate() - missing_args = self._missing_mapped_inputs_outputs() - if len(missing_args) > 0: - raise _user_exceptions.FlyteAssertion( - "The task {} is invalid because not all inputs and outputs in the " - "task function definition were specified in @outputs and @inputs. " - "We are missing definitions for {}.".format(self, missing_args) - ) - - @_exception_scopes.system_entry_point - def unit_test(self, **input_map): - """ - :param dict[Text, T] input_map: Python Std input from users. We will cast these to the appropriate Flyte - literals. - :returns: Depends on the behavior of the specific task in the unit engine. - """ - return ( - _engine_loader.get_engine("unit") - .get_task(self) - .execute( - _type_helpers.pack_python_std_map_to_literal_map( - input_map, - { - k: _type_helpers.get_sdk_type_from_literal_type(v.type) - for k, v in _six.iteritems(self.interface.inputs) - }, - ) - ) - ) - - @_exception_scopes.system_entry_point - def local_execute(self, **input_map): - """ - :param dict[Text, T] input_map: Python Std input from users. We will cast these to the appropriate Flyte - literals. - :rtype: dict[Text, T] - :returns: The output produced by this task in Python standard format. - """ - return ( - _engine_loader.get_engine("local") - .get_task(self) - .execute( - _type_helpers.pack_python_std_map_to_literal_map( - input_map, - { - k: _type_helpers.get_sdk_type_from_literal_type(v.type) - for k, v in _six.iteritems(self.interface.inputs) - }, - ) - ) - ) - - def _execute_user_code(self, context, inputs): - """ - :param flytekit.engines.common.EngineContext context: - :param dict[Text, T] inputs: This variable is a bit of a misnomer, since it's both inputs and outputs. The - dictionary passed here will be passed to the user-defined function, and will have values that are a - variety of types. The T's here are Python std values for inputs. If there isn't a native Python type for - something (like Schema or Blob), they are the Flyte classes. For outputs they are OutputReferences. - (Note that these are not the same OutputReferences as in BindingData's) - :rtype: Any: the returned object from user code. - :returns: This function must return a dictionary mapping 'filenames' to Flyte Interface Entities. These - entities will be used by the engine to pass data from node to node, populate metadata, etc. etc.. Each - engine will have different behavior. For instance, the Flyte engine will upload the entities to a remote - working directory (with the names provided), which will in turn allow Flyte Propeller to push along the - workflow. Where as local engine will merely feed the outputs directly into the next node. - """ - if self.task_style == SdkRunnableTaskStyle.V0: - return _exception_scopes.user_entry_point(self.task_function)( - ExecutionParameters( - execution_date=context.execution_date, - # TODO: it might be better to consider passing the full struct - execution_id=_six.text_type(WorkflowExecutionIdentifier.promote_from_model(context.execution_id)), - stats=context.stats, - logging=context.logging, - tmp_dir=context.working_directory, - ), - **inputs, - ) - - @_exception_scopes.system_entry_point - def execute(self, context, inputs): - """ - :param flytekit.engines.common.EngineContext context: - :param flytekit.models.literals.LiteralMap inputs: - :rtype: dict[Text, flytekit.models.common.FlyteIdlEntity] - :returns: This function must return a dictionary mapping 'filenames' to Flyte Interface Entities. These - entities will be used by the engine to pass data from node to node, populate metadata, etc. etc.. Each - engine will have different behavior. For instance, the Flyte engine will upload the entities to a remote - working directory (with the names provided), which will in turn allow Flyte Propeller to push along the - workflow. Where as local engine will merely feed the outputs directly into the next node. - """ - inputs_dict = _type_helpers.unpack_literal_map_to_sdk_python_std( - inputs, {k: _type_helpers.get_sdk_type_from_literal_type(v.type) for k, v in self.interface.inputs.items()} - ) - outputs_dict = { - name: _task_output.OutputReference(_type_helpers.get_sdk_type_from_literal_type(variable.type)) - for name, variable in _six.iteritems(self.interface.outputs) - } - - # Old style - V0: If annotations are used to define outputs, do not append outputs to the inputs dict - if not self.task_function.__annotations__ or "return" not in self.task_function.__annotations__: - inputs_dict.update(outputs_dict) - self._execute_user_code(context, inputs_dict) - return { - _constants.OUTPUT_FILE_NAME: _literal_models.LiteralMap( - literals={k: v.sdk_value for k, v in _six.iteritems(outputs_dict)} - ) - } - - @_exception_scopes.system_entry_point - def fast_register(self, project, domain, name, digest, additional_distribution, dest_dir) -> str: - """ - The fast register call essentially hijacks the task container commandline. - Say an existing task container definition had a commandline like so: - flyte_venv pyflyte-execute --task-module app.workflows.my_workflow --task-name my_task - - The fast register command introduces a wrapper call to fast-execute the original commandline like so: - flyte_venv pyflyte-fast-execute --additional-distribution s3://my-s3-bucket/foo/bar/12345.tar.gz -- - flyte_venv pyflyte-execute --task-module app.workflows.my_workflow --task-name my_task - - At execution time pyflyte-fast-execute will ensure the additional distribution (i.e. the fast-registered code) - exists before calling the original task commandline. - - :param Text project: The project in which to register this task. - :param Text domain: The domain in which to register this task. - :param Text name: The name to give this task. - :param Text digest: The version in which to register this task. - :param Text additional_distribution: User-specified location for remote source code distribution. - :param Text The optional location for where to install the additional distribution at runtime - :rtype: Text: Registered identifier. - """ - - original_container = self.container - container = _copy.deepcopy(original_container) - args = ["pyflyte-fast-execute", "--additional-distribution", additional_distribution] - if dest_dir: - args += ["--dest-dir", dest_dir] - args += ["--"] + container.args - container._args = args - self._container = container - - try: - registered_id = self.register(project, domain, name, digest) - except Exception: - self._container = original_container - raise - self._has_fast_registered = True - self._container = original_container - return str(registered_id) - - @property - def has_fast_registered(self) -> bool: - return self._has_fast_registered - - def _get_container_definition( - self, - storage_request=None, - cpu_request=None, - gpu_request=None, - memory_request=None, - storage_limit=None, - cpu_limit=None, - gpu_limit=None, - memory_limit=None, - environment=None, - cls=None, - ): - """ - :param Text storage_request: - :param Text cpu_request: - :param Text gpu_request: - :param Text memory_request: - :param Text storage_limit: - :param Text cpu_limit: - :param Text gpu_limit: - :param Text memory_limit: - :param dict[Text,Text] environment: - :param cls Optional[type]: Type of container to instantiate. Generally should subclass SdkRunnableContainer. - :rtype: flytekit.models.task.Container - """ - storage_limit = storage_limit or _resource_config.DEFAULT_STORAGE_LIMIT.get() - storage_request = storage_request or _resource_config.DEFAULT_STORAGE_REQUEST.get() - cpu_limit = cpu_limit or _resource_config.DEFAULT_CPU_LIMIT.get() - cpu_request = cpu_request or _resource_config.DEFAULT_CPU_REQUEST.get() - gpu_limit = gpu_limit or _resource_config.DEFAULT_GPU_LIMIT.get() - gpu_request = gpu_request or _resource_config.DEFAULT_GPU_REQUEST.get() - memory_limit = memory_limit or _resource_config.DEFAULT_MEMORY_LIMIT.get() - memory_request = memory_request or _resource_config.DEFAULT_MEMORY_REQUEST.get() - - resources = SdkRunnableContainer.get_resources( - storage_request, cpu_request, gpu_request, memory_request, storage_limit, cpu_limit, gpu_limit, memory_limit - ) - - return (cls or SdkRunnableContainer)( - command=[], - args=[ - "pyflyte-execute", - "--task-module", - self.task_module, - "--task-name", - self.task_function_name, - "--inputs", - "{{.input}}", - "--output-prefix", - "{{.outputPrefix}}", - "--raw-output-data-prefix", - "{{.rawOutputDataPrefix}}", - ], - resources=resources, - env=environment, - config={}, - ) - - def _validate_inputs(self, inputs): - """ - This method should be overridden in sub-classes that intend to do additional checks on inputs. If validation - fails, this function should raise an informative exception. - :param dict[Text, flytekit.models.interface.Variable] inputs: Input variables to validate - :raises: flytekit.common.exceptions.user.FlyteValidationException - """ - super(SdkRunnableTask, self)._validate_inputs(inputs) - for k, v in _six.iteritems(inputs): - if not self._is_argname_in_function_definition(k): - raise _user_exceptions.FlyteValidationException( - "The input named '{}' was not specified in the task function. Therefore, this input cannot be " - "provided to the task.".format(k) - ) - if _type_helpers.get_sdk_type_from_literal_type(v.type) in type(self)._banned_inputs: - raise _user_exceptions.FlyteValidationException( - "The input '{}' is not an accepted input type.".format(v) - ) - - def _validate_outputs(self, outputs): - """ - This method should be overridden in sub-classes that intend to do additional checks on outputs. If validation - fails, this function should raise an informative exception. - :param dict[Text, flytekit.models.interface.Variable] outputs: Output variables to validate - :raises: flytekit.common.exceptions.user.FlyteValidationException - """ - super(SdkRunnableTask, self)._validate_outputs(outputs) - for k, v in _six.iteritems(outputs): - if not self._is_argname_in_function_definition(k): - raise _user_exceptions.FlyteValidationException( - "The output named '{}' was not specified in the task function. Therefore, this output cannot be " - "provided to the task.".format(k) - ) - if _type_helpers.get_sdk_type_from_literal_type(v.type) in type(self)._banned_outputs: - raise _user_exceptions.FlyteValidationException( - "The output '{}' is not an accepted output type.".format(v) - ) - - def _get_kwarg_inputs(self): - # Trim off first parameter as it is reserved for workflow_parameters - return set(_getargspec(self.task_function).args[1:]) - - def _is_argname_in_function_definition(self, key): - return key in self._get_kwarg_inputs() - - def _missing_mapped_inputs_outputs(self): - # Trim off first parameter as it is reserved for workflow_parameters - args = self._get_kwarg_inputs() - inputs_and_outputs = set(self.interface.outputs.keys()) | set(self.interface.inputs.keys()) - return args ^ inputs_and_outputs diff --git a/flytekit/common/tasks/sidecar_task.py b/flytekit/common/tasks/sidecar_task.py deleted file mode 100644 index 15cb62d760..0000000000 --- a/flytekit/common/tasks/sidecar_task.py +++ /dev/null @@ -1,245 +0,0 @@ -import six as _six -from flyteidl.core import tasks_pb2 as _core_task -from google.protobuf.json_format import MessageToDict as _MessageToDict - -from flytekit.common import sdk_bases as _sdk_bases -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.tasks import sdk_dynamic as _sdk_dynamic -from flytekit.common.tasks import sdk_runnable as _sdk_runnable -from flytekit.models import task as _task_models -from flytekit.plugins import k8s as _lazy_k8s - - -class SdkSidecarTask(_sdk_runnable.SdkRunnableTask, metaclass=_sdk_bases.ExtendedSdkType): - - """ - This class includes the additional logic for building a task that executes as a Sidecar Job. - - """ - - def __init__( - self, - task_function, - task_type, - discovery_version, - retries, - interruptible, - deprecated, - storage_request, - cpu_request, - gpu_request, - memory_request, - storage_limit, - cpu_limit, - gpu_limit, - memory_limit, - discoverable, - timeout, - environment, - cache_serializable, - pod_spec=None, - primary_container_name=None, - annotations=None, - labels=None, - ): - """ - :param _sdk_runnable.SdkRunnableTask sdk_runnable_task: - :param generated_pb2.PodSpec pod_spec: - :param Text primary_container_name: - :param dict[Text, Text] annotations: - :param dict[Text, Text] labels: - :raises: flytekit.common.exceptions.user.FlyteValidationException - """ - if not pod_spec: - raise _user_exceptions.FlyteValidationException("A pod spec cannot be undefined") - if not primary_container_name: - raise _user_exceptions.FlyteValidationException("A primary container name cannot be undefined") - - super(SdkSidecarTask, self).__init__( - task_function, - task_type, - discovery_version, - retries, - interruptible, - deprecated, - storage_request, - cpu_request, - gpu_request, - memory_request, - storage_limit, - cpu_limit, - gpu_limit, - memory_limit, - discoverable, - timeout, - environment, - cache_serializable, - custom=None, - ) - - self.reconcile_partial_pod_spec_and_task(pod_spec, primary_container_name, annotations, labels) - - def reconcile_partial_pod_spec_and_task(self, pod_spec, primary_container_name, annotations=None, labels=None): - """ - Assigns the custom field as a the reconciled primary container and pod spec definition. - :param _sdk_runnable.SdkRunnableTask sdk_runnable_task: - :param generated_pb2.PodSpec pod_spec: - :param Text primary_container_name: - :param dict[Text, Text] annotations: - :param dict[Text, Text] labels: - :rtype: SdkSidecarTask - """ - - # First, insert a placeholder primary container if it is not defined in the pod spec. - containers = pod_spec.containers - primary_exists = False - for container in containers: - if container.name == primary_container_name: - primary_exists = True - break - if not primary_exists: - containers.extend([_lazy_k8s.io.api.core.v1.generated_pb2.Container(name=primary_container_name)]) - - final_containers = [] - for container in containers: - # In the case of the primary container, we overwrite specific container attributes with the default values - # used in an SDK runnable task. - if container.name == primary_container_name: - container.image = self._container.image - # clear existing commands - del container.command[:] - container.command.extend(self._container.command) - # also clear existing args - del container.args[:] - container.args.extend(self._container.args) - - resource_requirements = _lazy_k8s.io.api.core.v1.generated_pb2.ResourceRequirements() - for resource in self._container.resources.limits: - resource_requirements.limits[ - _core_task.Resources.ResourceName.Name(resource.name).lower() - ].CopyFrom(_lazy_k8s.io.apimachinery.pkg.api.resource.generated_pb2.Quantity(string=resource.value)) - for resource in self._container.resources.requests: - resource_requirements.requests[ - _core_task.Resources.ResourceName.Name(resource.name).lower() - ].CopyFrom(_lazy_k8s.io.apimachinery.pkg.api.resource.generated_pb2.Quantity(string=resource.value)) - if resource_requirements.ByteSize(): - # Important! Only copy over resource requirements if they are non-empty. - container.resources.CopyFrom(resource_requirements) - - del container.env[:] - container.env.extend( - [ - _lazy_k8s.io.api.core.v1.generated_pb2.EnvVar(name=key, value=val) - for key, val in _six.iteritems(self._container.env) - ] - ) - - final_containers.append(container) - - del pod_spec.containers[:] - pod_spec.containers.extend(final_containers) - - sidecar_job_plugin = _task_models.SidecarJob( - pod_spec=pod_spec, - primary_container_name=primary_container_name, - annotations=annotations, - labels=labels, - ).to_flyte_idl() - - self.assign_custom_and_return(_MessageToDict(sidecar_job_plugin)) - - -class SdkDynamicSidecarTask( - _sdk_dynamic.SdkDynamicTaskMixin, - SdkSidecarTask, - metaclass=_sdk_bases.ExtendedSdkType, -): - - """ - This class includes the additional logic for building a task that runs as - a Sidecar Job and executes parent-child tasks. - - """ - - def __init__( - self, - task_function, - task_type, - discovery_version, - retries, - interruptible, - deprecated, - storage_request, - cpu_request, - gpu_request, - memory_request, - storage_limit, - cpu_limit, - gpu_limit, - memory_limit, - discoverable, - timeout, - allowed_failure_ratio, - max_concurrency, - environment, - cache_serializable, - pod_spec=None, - primary_container_name=None, - annotations=None, - labels=None, - ): - """ - :param task_function: Function container user code. This will be executed via the SDK's engine. - :param Text task_type: string describing the task type - :param Text discovery_version: string describing the version for task discovery purposes - :param int retries: Number of retries to attempt - :param bool interruptible: Whether or not task is interruptible - :param Text deprecated: - :param Text storage_request: - :param Text cpu_request: - :param Text gpu_request: - :param Text memory_request: - :param Text storage_limit: - :param Text cpu_limit: - :param Text gpu_limit: - :param Text memory_limit: - :param bool discoverable: - :param datetime.timedelta timeout: - :param float allowed_failure_ratio: - :param int max_concurrency: - :param dict[Text, Text] environment: - :param bool cache_serializable: - :param generated_pb2.PodSpec pod_spec: - :param Text primary_container_name: - :param dict[Text, Text] annotations: - :param dict[Text, Text] labels: - :raises: flytekit.common.exceptions.user.FlyteValidationException - """ - - SdkSidecarTask.__init__( - self, - task_function, - task_type, - discovery_version, - retries, - interruptible, - deprecated, - storage_request, - cpu_request, - gpu_request, - memory_request, - storage_limit, - cpu_limit, - gpu_limit, - memory_limit, - discoverable, - timeout, - environment, - cache_serializable, - pod_spec=pod_spec, - primary_container_name=primary_container_name, - annotations=annotations, - labels=labels, - ) - - _sdk_dynamic.SdkDynamicTaskMixin.__init__(self, allowed_failure_ratio, max_concurrency) diff --git a/flytekit/common/tasks/spark_task.py b/flytekit/common/tasks/spark_task.py deleted file mode 100644 index f3f55f211e..0000000000 --- a/flytekit/common/tasks/spark_task.py +++ /dev/null @@ -1,208 +0,0 @@ -import typing - -try: - from inspect import getfullargspec as _getargspec -except ImportError: - from inspect import getargspec as _getargspec - -import copy as _copy -import hashlib as _hashlib -import json as _json -import os as _os -import sys as _sys - -import six as _six -from google.protobuf.json_format import MessageToDict as _MessageToDict - -from flytekit.bin import entrypoint as _entrypoint -from flytekit.common import constants as _constants -from flytekit.common.exceptions import scopes as _exception_scopes -from flytekit.common.tasks import output as _task_output -from flytekit.common.tasks import sdk_runnable as _sdk_runnable -from flytekit.common.types import helpers as _type_helpers -from flytekit.models import literals as _literal_models -from flytekit.models import task as _task_models -from flytekit.plugins import pyspark as _pyspark - - -class GlobalSparkContext(object): - _SPARK_CONTEXT = None - _SPARK_SESSION = None - - @classmethod - def get_spark_context(cls): - return cls._SPARK_CONTEXT - - @classmethod - def get_spark_session(cls): - return cls._SPARK_SESSION - - def __enter__(self): - GlobalSparkContext._SPARK_CONTEXT = _pyspark.SparkContext() - GlobalSparkContext._SPARK_SESSION = _pyspark.sql.SparkSession.builder.appName( - "Flyte Spark SQL Context" - ).getOrCreate() - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - GlobalSparkContext._SPARK_CONTEXT.stop() - GlobalSparkContext._SPARK_CONTEXT = None - return False - - -class SdkRunnableSparkContainer(_sdk_runnable.SdkRunnableContainer): - @property - def args(self): - """ - Override args to remove the injection of command prefixes - :rtype: list[Text] - """ - return self._args - - -class SdkSparkTask(_sdk_runnable.SdkRunnableTask): - """ - This class includes the additional logic for building a task that executes as a Spark Job. - - """ - - def __init__( - self, - task_function, - task_type, - discovery_version, - retries, - interruptible, - deprecated, - discoverable, - timeout, - spark_type, - spark_conf, - hadoop_conf, - environment, - cache_serializable, - ): - """ - :param task_function: Function container user code. This will be executed via the SDK's engine. - :param Text task_type: string describing the task type - :param Text discovery_version: string describing the version for task discovery purposes - :param int retries: Number of retries to attempt - :param bool interruptible: Whether or not task is interruptible - :param Text deprecated: - :param bool discoverable: - :param datetime.timedelta timeout: - :param dict[Text,Text] spark_conf: - :param dict[Text,Text] hadoop_conf: - :param dict[Text,Text] environment: [optional] environment variables to set when executing this task. - :param bool cache_serializable: - """ - - spark_exec_path = _os.path.abspath(_entrypoint.__file__) - if spark_exec_path.endswith(".pyc"): - spark_exec_path = spark_exec_path[:-1] - - self._spark_job = _task_models.SparkJob( - spark_conf=spark_conf, - hadoop_conf=hadoop_conf, - application_file="local://" + spark_exec_path, - executor_path=_sys.executable, - main_class="", - spark_type=spark_type, - ) - super(SdkSparkTask, self).__init__( - task_function, - task_type, - discovery_version, - retries, - interruptible, - deprecated, - "", - "", - "", - "", - "", - "", - "", - "", - discoverable, - timeout, - environment, - cache_serializable, - _MessageToDict(self._spark_job.to_flyte_idl()), - ) - - @_exception_scopes.system_entry_point - def execute(self, context, inputs): - """ - :param flytekit.engines.common.EngineContext context: - :param flytekit.models.literals.LiteralMap inputs: - :rtype: dict[Text,flytekit.models.common.FlyteIdlEntity] - :returns: This function must return a dictionary mapping 'filenames' to Flyte Interface Entities. These - entities will be used by the engine to pass data from node to node, populate metadata, etc. etc.. Each - engine will have different behavior. For instance, the Flyte engine will upload the entities to a remote - working directory (with the names provided), which will in turn allow Flyte Propeller to push along the - workflow. Where as local engine will merely feed the outputs directly into the next node. - """ - inputs_dict = _type_helpers.unpack_literal_map_to_sdk_python_std( - inputs, - {k: _type_helpers.get_sdk_type_from_literal_type(v.type) for k, v in _six.iteritems(self.interface.inputs)}, - ) - outputs_dict = { - name: _task_output.OutputReference(_type_helpers.get_sdk_type_from_literal_type(variable.type)) - for name, variable in _six.iteritems(self.interface.outputs) - } - - inputs_dict.update(outputs_dict) - - with GlobalSparkContext(): - _exception_scopes.user_entry_point(self.task_function)( - _sdk_runnable.ExecutionParameters( - execution_date=context.execution_date, - tmp_dir=context.working_directory, - stats=context.stats, - execution_id=context.execution_id, - logging=context.logging, - ), - GlobalSparkContext.get_spark_context(), - **inputs_dict - ) - return { - _constants.OUTPUT_FILE_NAME: _literal_models.LiteralMap( - literals={k: v.sdk_value for k, v in _six.iteritems(outputs_dict)} - ) - } - - @property - def spark_conf(self): - return self._spark_job.spark_conf - - @property - def hadoop_conf(self): - return self._spark_job.hadoop_conf - - def _get_container_definition(self, **kwargs): - """ - :rtype: SdkRunnableSparkContainer - """ - return super(SdkSparkTask, self)._get_container_definition(cls=SdkRunnableSparkContainer, **kwargs) - - def _get_kwarg_inputs(self): - # Trim off first two parameters as they are reserved for workflow_parameters and spark_context - return set(_getargspec(self.task_function).args[2:]) - - def with_overrides( - self, new_spark_conf: typing.Dict[str, str] = None, new_hadoop_conf: typing.Dict[str, str] = None - ): - """ - Creates a new SparkJob instance with the modified configuration or timeouts - """ - tk = _copy.deepcopy(self) - tk._spark_job = self._spark_job.with_overrides(new_spark_conf, new_hadoop_conf) - tk._custom = _MessageToDict(tk._spark_job.to_flyte_idl()) - - salt = _hashlib.md5(_json.dumps(tk.custom, sort_keys=True).encode("utf-8")).hexdigest() - tk._id._name = "{}-{}".format(self._id.name, salt) - # We are overriding the platform name creation to prevent problems in dynamic - tk.assign_name(tk._id._name) - - return tk diff --git a/flytekit/common/tasks/task.py b/flytekit/common/tasks/task.py deleted file mode 100644 index ba55399382..0000000000 --- a/flytekit/common/tasks/task.py +++ /dev/null @@ -1,423 +0,0 @@ -import hashlib as _hashlib -import json as _json -import logging as _logging -import uuid as _uuid - -import six as _six -from google.protobuf import json_format as _json_format -from google.protobuf import struct_pb2 as _struct - -from flytekit.common import interface as _interfaces -from flytekit.common import nodes as _nodes -from flytekit.common import sdk_bases as _sdk_bases -from flytekit.common import workflow_execution as _workflow_execution -from flytekit.common.core import identifier as _identifier -from flytekit.common.exceptions import scopes as _exception_scopes -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.mixins import hash as _hash_mixin -from flytekit.common.mixins import launchable as _launchable_mixin -from flytekit.common.mixins import registerable as _registerable -from flytekit.common.types import helpers as _type_helpers -from flytekit.configuration import auth as _auth_config -from flytekit.configuration import internal as _internal_config -from flytekit.configuration import sdk as _sdk_config -from flytekit.engines.flyte import engine as _flyte_engine -from flytekit.models import common as _common_model -from flytekit.models import execution as _admin_execution_models -from flytekit.models import task as _task_model -from flytekit.models.admin import common as _admin_common -from flytekit.models.core import identifier as _identifier_model -from flytekit.models.core import workflow as _workflow_model - - -class SdkTask( - _hash_mixin.HashOnReferenceMixin, - _registerable.RegisterableEntity, - _launchable_mixin.LaunchableEntity, - _task_model.TaskTemplate, - metaclass=_sdk_bases.ExtendedSdkType, -): - def __init__( - self, type, metadata, interface, custom, container=None, task_type_version=0, security_context=None, config=None - ): - """ - :param Text type: This is used to define additional extensions for use by Propeller or SDK. - :param TaskMetadata metadata: This contains information needed at runtime to determine behavior such as - whether or not outputs are discoverable, timeouts, and retries. - :param flytekit.common.interface.TypedInterface interface: The interface definition for this task. - :param dict[Text, T] custom: Arbitrary type for use by plugins. - :param Container container: Provides the necessary entrypoint information for execution. For instance, - a Container might be specified with the necessary command line arguments. - :param int task_type_version: Specific version of this task type used by plugins to potentially modify - execution behavior or serialization. - :param _SecurityContext security_context: - """ - # TODO: Remove the identifier portion and fill in with local values. - super(SdkTask, self).__init__( - _identifier.Identifier( - _identifier_model.ResourceType.TASK, - _internal_config.PROJECT.get(), - _internal_config.DOMAIN.get(), - _uuid.uuid4().hex, - _internal_config.VERSION.get(), - ), - type, - metadata, - interface, - custom, - container=container, - task_type_version=task_type_version, - security_context=security_context, - config=config, - ) - - @property - def interface(self): - """ - :rtype: flytekit.common.interface.TypedInterface - """ - return super(SdkTask, self).interface - - @property - def resource_type(self): - """ - Integer from _identifier.ResourceType enum - :rtype: int - """ - return _identifier_model.ResourceType.TASK - - @property - def entity_type_text(self): - """ - :rtype: Text - """ - return "Task" - - @classmethod - def promote_from_model(cls, base_model): - """ - :param flytekit.models.task.TaskTemplate base_model: - :rtype: SdkTask - """ - t = cls( - type=base_model.type, - metadata=base_model.metadata, - interface=_interfaces.TypedInterface.promote_from_model(base_model.interface), - custom=base_model.custom, - container=base_model.container, - task_type_version=base_model.task_type_version, - ) - # Override the newly generated name if one exists in the base model - if not base_model.id.is_empty: - t._id = _identifier.Identifier.promote_from_model(base_model.id) - - return t - - def assign_custom_and_return(self, custom): - self._custom = custom - return self - - def assign_type_and_return(self, new_type): - self._type = new_type - return self - - @_exception_scopes.system_entry_point - def __call__(self, *args, **input_map): - """ - :param list[T] args: Do not specify. Kwargs only are supported for this function. - :param dict[str, T] input_map: Map of inputs. Can be statically defined or OutputReference links. - :rtype: flytekit.common.nodes.SdkNode - """ - if len(args) > 0: - raise _user_exceptions.FlyteAssertion( - "When adding a task as a node in a workflow, all inputs must be specified with kwargs only. We " - "detected {} positional args.".format(len(args)) - ) - - bindings, upstream_nodes = self.interface.create_bindings_for_inputs(input_map) - - # TODO: Remove DEADBEEF - # One thing to note - this function is not overloaded at the SdkRunnableTask layer, which means 'self' here - # will sometimes refer to an object that can be executed locally, and other times will refer to something - # that cannot (ie a pure SdkTask object, fetched from Admin for instance). - return _nodes.SdkNode( - id=None, - metadata=_workflow_model.NodeMetadata( - "DEADBEEF", - self.metadata.timeout, - self.metadata.retries, - self.metadata.interruptible, - ), - bindings=sorted(bindings, key=lambda b: b.var), - upstream_nodes=upstream_nodes, - sdk_task=self, - ) - - @_exception_scopes.system_entry_point - def register(self, project, domain, name, version): - """ - :param Text project: The project in which to register this task. - :param Text domain: The domain in which to register this task. - :param Text name: The name to give this task. - :param Text version: The version in which to register this task. - """ - # TODO: Revisit the notion of supplying the project, domain, name, version, as opposed to relying on the - # current ID. - self.validate() - id_to_register = _identifier.Identifier(_identifier_model.ResourceType.TASK, project, domain, name, version) - old_id = self.id - - client = _flyte_engine.get_client() - try: - self._id = id_to_register - client.create_task(id_to_register, _task_model.TaskSpec(self)) - self._id = old_id - self._has_registered = True - return str(id_to_register) - except _user_exceptions.FlyteEntityAlreadyExistsException: - pass - except Exception: - self._id = old_id - raise - - @_exception_scopes.system_entry_point - def serialize(self): - """ - :rtype: flyteidl.admin.task_pb2.TaskSpec - """ - return _task_model.TaskSpec(self).to_flyte_idl() - - @classmethod - @_exception_scopes.system_entry_point - def fetch(cls, project, domain, name, version): - """ - This function uses the engine loader to call create a hydrated task from Admin. - :param Text project: - :param Text domain: - :param Text name: - :param Text version: - :rtype: SdkTask - """ - task_id = _identifier.Identifier(_identifier_model.ResourceType.TASK, project, domain, name, version) - admin_task = _flyte_engine.get_client().get_task(task_id) - - sdk_task = cls.promote_from_model(admin_task.closure.compiled_task.template) - sdk_task._id = task_id - sdk_task._has_registered = True - return sdk_task - - @classmethod - @_exception_scopes.system_entry_point - def fetch_latest(cls, project, domain, name): - """ - This function uses the engine loader to call create a latest hydrated task from Admin. - :param Text project: - :param Text domain: - :param Text name: - :rtype: SdkTask - """ - named_task = _common_model.NamedEntityIdentifier(project, domain, name) - client = _flyte_engine.get_client() - task_list, _ = client.list_tasks_paginated( - named_task, - limit=1, - sort_by=_admin_common.Sort("created_at", _admin_common.Sort.Direction.DESCENDING), - ) - admin_task = task_list[0] if task_list else None - - if not admin_task: - raise _user_exceptions.FlyteEntityNotExistException("Named task {} not found".format(named_task)) - sdk_task = cls.promote_from_model(admin_task.closure.compiled_task.template) - sdk_task._id = admin_task.id - return sdk_task - - @_exception_scopes.system_entry_point - def validate(self): - pass - - @_exception_scopes.system_entry_point - def add_inputs(self, inputs): - raise _user_exceptions.FlyteUserException("You can not add inputs to this task") - - @_exception_scopes.system_entry_point - def add_outputs(self, outputs): - """ - Adds the outputs to this task. This can be called multiple times, but it will fail if an output with a given - name is added more than once, a name collides with an input, or if the name doesn't exist as an arg name in - the wrapped function. - :param dict[Text, flytekit.models.interface.Variable] outputs: names and variables to add as outputs - to this task - """ - self._validate_outputs(outputs) - self.interface.outputs.update(outputs) - - def _validate_inputs(self, inputs): - """ - This method should be overridden in sub-classes that intend to do additional checks on inputs. If validation - fails, this function should raise an informative exception. - :param dict[Text, flytekit.models.interface.Variable] inputs: Input variables to validate - :raises: flytekit.common.exceptions.user.FlyteValidationException - """ - for k, v in _six.iteritems(inputs): - if k in self.interface.inputs: - raise _user_exceptions.FlyteValidationException( - "An input with name '{}' is already defined. Redefinition is not allowed.".format(k) - ) - if k in self.interface.outputs: - raise _user_exceptions.FlyteValidationException( - "An output with name '{}' is already defined. Therefore '{}' can't be defined as an " - "input".format(k, v) - ) - - def _validate_outputs(self, outputs): - """ - This method should be overridden in sub-classes that intend to do additional checks on outputs. If validation - fails, this function should raise an informative exception. - :param dict[Text, flytekit.models.interface.Variable] outputs: Output variables to validate - :raises: flytekit.common.exceptions.user.FlyteValidationException - """ - for k, v in _six.iteritems(outputs): - if k in self.interface.outputs: - raise _user_exceptions.FlyteValidationException( - "An output with name '{}' is already defined. Redefinition is not allowed.".format(k) - ) - if k in self.interface.inputs: - raise _user_exceptions.FlyteValidationException( - "An input with name '{}' is already defined. Therefore '{}' can't be defined as an " - "input".format(k, v) - ) - - def __repr__(self): - return "Flyte {task_type}: {interface}".format(task_type=self.type, interface=self.interface) - - def _python_std_input_map_to_literal_map(self, inputs): - """ - :param dict[Text,Any] inputs: A dictionary of Python standard inputs that will be type-checked and compiled - to a LiteralMap - :rtype: flytekit.models.literals.LiteralMap - """ - return _type_helpers.pack_python_std_map_to_literal_map( - inputs, - {k: _type_helpers.get_sdk_type_from_literal_type(v.type) for k, v in _six.iteritems(self.interface.inputs)}, - ) - - def _produce_deterministic_version(self, version=None): - """ - :param Text version: - :return Text: - """ - - if self.container is not None and self.container.data_loading_config is None: - # Only in the case of raw container tasks (which are the only valid tasks with container definitions that - # can assign a client-side task version) their data config will be None. - raise ValueError("Client-side task versions are not supported for {} task type".format(self.type)) - if version is not None: - return version - custom = _json_format.Parse(_json.dumps(self.custom, sort_keys=True), _struct.Struct()) if self.custom else None - - # The task body is the entirety of the task template MINUS the identifier. The identifier is omitted because - # 1) this method is used to compute the version portion of the identifier and - # 2 ) the SDK will actually generate a unique name on every task instantiation which is not great for - # the reproducibility this method attempts. - task_body = ( - self.type, - self.metadata.to_flyte_idl().SerializeToString(deterministic=True), - self.interface.to_flyte_idl().SerializeToString(deterministic=True), - custom, - ) - return _hashlib.md5(str(task_body).encode("utf-8")).hexdigest() - - @_exception_scopes.system_entry_point - def register_and_launch(self, project, domain, name, version=None, inputs=None): - """ - :param Text project: The project in which to register and launch this task. - :param Text domain: The domain in which to register and launch this task. - :param Text name: The name to give this task. - :param Text version: The version in which to register this task - :param dict[Text, Any] inputs: A dictionary of Python standard inputs that will be type-checked, then compiled - to a LiteralMap. - - :rtype: flytekit.common.workflow_execution.SdkWorkflowExecution - """ - self.validate() - version = self._produce_deterministic_version(version) - self.register(project, domain, name, version) - return self.launch(project, domain, inputs=inputs) - - @_exception_scopes.system_entry_point - def launch_with_literals( - self, - project, - domain, - literal_inputs, - name=None, - notification_overrides=None, - label_overrides=None, - annotation_overrides=None, - auth_role=None, - ): - """ - Launches a single task execution and returns the execution identifier. - :param Text project: - :param Text domain: - :param flytekit.models.literals.LiteralMap literal_inputs: Inputs to the execution. - :param Text name: [Optional] If specified, an execution will be created with this name. Note: the name must - be unique within the context of the project and domain. - :param list[flytekit.common.notifications.Notification] notification_overrides: [Optional] If specified, these - are the notifications that will be honored for this execution. An empty list signals to disable all - notifications. - :param flytekit.models.common.Labels label_overrides: - :param flytekit.models.common.Annotations annotation_overrides: - :param flytekit.models.common.AuthRole auth_role: - :rtype: flytekit.common.workflow_execution.SdkWorkflowExecution - """ - disable_all = notification_overrides == [] - if disable_all: - notification_overrides = None - else: - notification_overrides = _admin_execution_models.NotificationList(notification_overrides or []) - disable_all = None - - # Unlike regular workflow executions, single task executions must always specify an auth role, since there isn't - # any existing launch plan with a bound auth role to fall back on. - if auth_role is None: - assumable_iam_role = _auth_config.ASSUMABLE_IAM_ROLE.get() - kubernetes_service_account = _auth_config.KUBERNETES_SERVICE_ACCOUNT.get() - - if not (assumable_iam_role or kubernetes_service_account): - _logging.warning( - "Using deprecated `role` from config. " - "Please update your config to use `assumable_iam_role` instead" - ) - assumable_iam_role = _sdk_config.ROLE.get() - auth_role = _common_model.AuthRole( - assumable_iam_role=assumable_iam_role, - kubernetes_service_account=kubernetes_service_account, - ) - - client = _flyte_engine.get_client() - try: - # TODO(katrogan): Add handling to register the underlying task if it's not already. - exec_id = client.create_execution( - project, - domain, - name, - _admin_execution_models.ExecutionSpec( - self.id, - _admin_execution_models.ExecutionMetadata( - _admin_execution_models.ExecutionMetadata.ExecutionMode.MANUAL, - "sdk", # TODO: get principle - 0, # TODO: Detect nesting - ), - notifications=notification_overrides, - disable_all=disable_all, - labels=label_overrides, - annotations=annotation_overrides, - auth_role=auth_role, - ), - literal_inputs, - ) - except _user_exceptions.FlyteEntityAlreadyExistsException: - exec_id = _identifier.WorkflowExecutionIdentifier(project, domain, name) - execution = client.get_execution(exec_id) - return _workflow_execution.SdkWorkflowExecution.promote_from_model(execution) diff --git a/flytekit/common/types/__init__.py b/flytekit/common/types/__init__.py deleted file mode 100644 index b2e2a3729a..0000000000 --- a/flytekit/common/types/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -""" -This package contains the runtime logic wrapping our type models. -""" diff --git a/flytekit/common/types/base_sdk_types.py b/flytekit/common/types/base_sdk_types.py deleted file mode 100644 index ab12969865..0000000000 --- a/flytekit/common/types/base_sdk_types.py +++ /dev/null @@ -1,141 +0,0 @@ -import abc as _abc - -from flyteidl.core.literals_pb2 import Literal - -from flytekit.common import sdk_bases as _sdk_bases -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.models import common as _common_models -from flytekit.models import literals as _literal_models - - -class FlyteSdkType(_sdk_bases.ExtendedSdkType, metaclass=_common_models.FlyteABCMeta): - @_abc.abstractmethod - def is_castable_from(cls, other): - """ - :param flytekit.common.types.base_literal_types.FlyteSdkType other: - :rtype: bool - """ - pass - - @_abc.abstractmethod - def from_python_std(cls, t_value): - """ - :param T t_value: It is up to each individual object as to whether or not this value can be cast. - :rtype: FlyteSdkValue - :raises: flytekit.common.exceptions.user.FlyteTypeException - """ - pass - - @_abc.abstractmethod - def from_string(cls, string_value): - """ - :param Text string_value: It is up to each individual object to implement this. - :rtype: FlyteSdkValue - :raises: flytekit.common.exceptions.user.FlyteTypeException - """ - pass - - @_abc.abstractmethod - def promote_from_model(cls, literal): - """ - :param flytekit.models.literals.Literal literal: - :rtype: FlyteSdkValue - """ - pass - - @_abc.abstractmethod - def to_flyte_literal_type(cls): - """ - :rtype: flytekit.models.types.LiteralType - """ - pass - - def __hash__(cls): - return hash(cls.to_flyte_literal_type()) - - -class FlyteSdkValue(_literal_models.Literal, metaclass=FlyteSdkType): - @classmethod - def from_flyte_idl(cls, pb2_object: Literal): - """ - :param flyteidl.core.literals_pb2.Literal pb2_object: - :rtype: FlyteSdkValue - """ - literal = _literal_models.Literal.from_flyte_idl(pb2_object) - if literal.scalar is not None and literal.scalar.none_type is not None: - return Void() - return cls.promote_from_model(literal) - - @_abc.abstractmethod - def to_python_std(self): - pass - - -class InstantiableType(FlyteSdkType, metaclass=_common_models.FlyteABCMeta): - @_abc.abstractmethod - def __call__(cls, *args, **kwargs): - """ - TODO: Figure out generics for type hinting. - - :rtype: T - """ - return super(InstantiableType, cls).__call__(*args, **kwargs) - - -class Void(FlyteSdkValue): - @classmethod - def is_castable_from(cls, other): - """ - :param flytekit.common.types.base_literal_types.FlyteSdkType other: - :rtype: bool - """ - return True - - @classmethod - def from_python_std(cls, t_value): - """ - :param T t_value: It is up to each individual object as to whether or not this value can be cast. - :rtype: FlyteSdkValue - :raises: flytekit.common.exceptions.user.FlyteTypeException - """ - return cls() - - @classmethod - def to_flyte_literal_type(cls): - """ - :rtype: flytekit.models.types.LiteralType - """ - raise _user_exceptions.FlyteAssertion( - "A Void type does not have a literal type and cannot be used in this " "manner." - ) - - @classmethod - def promote_from_model(cls, _): - """ - Creates an object of this type from the model primitive defining it. - :param flytekit.models.literals.Literal _: - :rtype: Void - """ - return cls() - - @classmethod - def short_class_string(cls): - """ - :rtype: Text - """ - return "Void" - - def __init__(self): - super(Void, self).__init__(scalar=_literal_models.Scalar(none_type=_literal_models.Void())) - - def to_python_std(self): - """ - :rtype: NoneType - """ - return None - - def short_string(self): - """ - :rtype: Text - """ - return "Void()" diff --git a/flytekit/common/types/blobs.py b/flytekit/common/types/blobs.py deleted file mode 100644 index 7870cb75bd..0000000000 --- a/flytekit/common/types/blobs.py +++ /dev/null @@ -1,465 +0,0 @@ -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.types import base_sdk_types as _base_sdk_types -from flytekit.common.types.impl import blobs as _blob_impl -from flytekit.models import literals as _literals -from flytekit.models import types as _idl_types -from flytekit.models.core import types as _core_types - - -class BlobInstantiator(_base_sdk_types.InstantiableType): - @staticmethod - def create_at_known_location(location): - """ - :param Text location: - :rtype: flytekit.common.types.impl.blobs.Blob - """ - return _blob_impl.Blob.create_at_known_location(location, mode="wb") - - @staticmethod - def fetch(remote_path, local_path=None): - """ - :param Text remote_path: - :param Text local_path: [Optional] If specified, the Blob is copied to this location. If specified, - this location is NOT managed and the blob will not be cleaned up upon exit. - :rtype: flytekit.common.types.impl.blobs.Blob - """ - return _blob_impl.Blob.fetch(remote_path, mode="rb", local_path=local_path) - - def __call__(cls, *args, **kwargs): - """ - TODO: Is there a better way to deal with this? - - We want the behavior of Types.Blob() returns a _blob_impl.Blob, but also to be able to use this object to - wrap a _blob_impl.Blob via Types.Blob(_blob_impl.Blob()) for serialization, type checking, etc.. - - :rtype: flytekit.common.types.impl.blobs.Blob - """ - if not args and not kwargs: - return _blob_impl.Blob.create_at_any_location(mode="wb") - else: - return super(BlobInstantiator, cls).__call__(*args, **kwargs) - - -# TODO: Make blobs and schemas pluggable -class Blob(_base_sdk_types.FlyteSdkValue, metaclass=BlobInstantiator): - @classmethod - def from_string(cls, string_value): - """ - :param Text string_value: - :rtype: Blob - """ - if not string_value: - _user_exceptions.FlyteValueException(string_value, "Cannot create a Blob from the provided path value.") - return cls(_blob_impl.Blob.from_string(string_value, mode="rb")) - - @classmethod - def is_castable_from(cls, other): - """ - :param flytekit.common.types.base_literal_types.FlyteSdkType other: - :rtype: bool - """ - return cls == other - - @classmethod - def from_python_std(cls, t_value): - """ - :param T t_value: It is up to each individual object as to whether or not this value can be cast. - :rtype: FlyteSdkValue - :raises: flytekit.common.exceptions.user.FlyteTypeException - """ - if t_value is None: - return _base_sdk_types.Void() - elif isinstance(t_value, _blob_impl.Blob): - blob = t_value - else: - blob = _blob_impl.Blob.from_python_std(t_value) - return cls(blob) - - @classmethod - def to_flyte_literal_type(cls): - """ - :rtype: flytekit.models.types.LiteralType - """ - return _idl_types.LiteralType( - blob=_core_types.BlobType(format="", dimensionality=_core_types.BlobType.BlobDimensionality.SINGLE) - ) - - @classmethod - def promote_from_model(cls, literal_model): - """ - Creates an object of this type from the model primitive defining it. - :param flytekit.models.literals.Literal literal_model: - :rtype: Blob - """ - return cls(_blob_impl.Blob.promote_from_model(literal_model.scalar.blob)) - - @classmethod - def short_class_string(cls): - """ - :rtype: Text - """ - return "Blob" - - def __init__(self, value): - """ - :param flytekit.common.types.impl.blobs.Blob value: Blob value to wrap - """ - super(Blob, self).__init__(scalar=_literals.Scalar(blob=value)) - - def to_python_std(self): - """ - :rtype: flytekit.common.types.impl.blobs.Blob - """ - return self.scalar.blob - - def short_string(self): - """ - :rtype: Text - """ - return "Blob(uri={}{})".format( - self.scalar.blob.uri, - ", format={}".format(self.scalar.blob.metadata.type.format) - if self.scalar.blob.metadata.type.format - else "", - ) - - -class MultiPartBlobInstantiator(_base_sdk_types.InstantiableType): - @staticmethod - def create_at_known_location(location): - """ - :param Text location: - :rtype: flytekit.common.types.impl.blobs.MultiPartBlob - """ - return _blob_impl.MultiPartBlob.create_at_known_location(location, mode="wb") - - @staticmethod - def fetch(remote_path, local_path=None): - """ - :param Text remote_path: - :param Text local_path: [Optional] If specified, the MultiPartBlob is copied to this location. If specified, - this location is NOT managed and the blob will not be cleaned up upon exit. - :rtype: flytekit.common.types.impl.blobs.MultiPartBlob - """ - return _blob_impl.MultiPartBlob.fetch(remote_path, mode="rb", local_path=local_path) - - def __call__(cls, *args, **kwargs): - """ - TODO: Is there a better way to deal with this? - - We want the behavior of Types.MultiPartBlob() returns a _blob_impl.MultiPartBlob, but also to be able to use - this object to wrap a _blob_impl.MultiPartBlob via Types.MultiPartBlob(_blob_impl.MultiPartBlob()) for - serialization, type checking, etc.. - - :rtype: flytekit.common.types.impl.blobs.MultiPartBlob - """ - if not args and not kwargs: - return _blob_impl.MultiPartBlob.create_at_any_location(mode="wb") - else: - return super(MultiPartBlobInstantiator, cls).__call__(*args, **kwargs) - - -class MultiPartBlob(_base_sdk_types.FlyteSdkValue, metaclass=MultiPartBlobInstantiator): - @classmethod - def from_string(cls, string_value): - """ - :param Text string_value: - :rtype: MultiPartBlob - """ - if not string_value: - _user_exceptions.FlyteValueException( - string_value, - "Cannot create a MultiPartBlob from the provided path " "value.", - ) - return cls(_blob_impl.MultiPartBlob.from_string(string_value, mode="rb")) - - @classmethod - def is_castable_from(cls, other): - """ - :param flytekit.common.types.base_literal_types.FlyteSdkType other: - :rtype: bool - """ - return cls == other - - @classmethod - def from_python_std(cls, t_value): - """ - :param T t_value: It is up to each individual object as to whether or not this value can be cast. - :rtype: FlyteSdkValue - :raises: flytekit.common.exceptions.user.FlyteTypeException - """ - if t_value is None: - return _base_sdk_types.Void() - elif isinstance(t_value, _blob_impl.MultiPartBlob): - blob = t_value - else: - blob = _blob_impl.MultiPartBlob.from_python_std(t_value) - return cls(blob) - - @classmethod - def to_flyte_literal_type(cls): - """ - :rtype: flytekit.models.types.LiteralType - """ - return _idl_types.LiteralType( - blob=_core_types.BlobType( - format="", - dimensionality=_core_types.BlobType.BlobDimensionality.MULTIPART, - ) - ) - - @classmethod - def promote_from_model(cls, literal_model): - """ - Creates an object of this type from the model primitive defining it. - :param flytekit.models.literals.Literal literal_model: - :rtype: MultiPartBlob - """ - return cls(_blob_impl.MultiPartBlob.promote_from_model(literal_model.scalar.blob)) - - @classmethod - def short_class_string(cls): - """ - :rtype: Text - """ - return "MultiPartBlob" - - def __init__(self, value): - """ - :param flytekit.common.types.impl.blobs.MultiPartBlob value: Blob value to wrap - """ - super(MultiPartBlob, self).__init__(scalar=_literals.Scalar(blob=value)) - - def to_python_std(self): - """ - :rtype: flytekit.common.types.impl.blobs.MultiPartBlob - """ - return self.scalar.blob - - def short_string(self): - """ - :rtype: Text - """ - return "MultiPartBlob(uri={}{})".format( - self.scalar.blob.uri, - ", format={}".format(self.scalar.blob.metadata.type.format) - if self.scalar.blob.metadata.type.format - else "", - ) - - -class CsvInstantiator(BlobInstantiator): - @staticmethod - def create_at_known_location(location): - """ - :param Text location: - :rtype: flytekit.common.types.impl.blobs.CSV - """ - return _blob_impl.Blob.create_at_known_location(location, mode="w", format="csv") - - @staticmethod - def fetch(remote_path, local_path=None): - """ - :param Text remote_path: - :param Text local_path: [Optional] If specified, the MultiPartBlob is copied to this location. If specified, - this location is NOT managed and the blob will not be cleaned up upon exit. - :rtype: flytekit.common.types.impl.blobs.CSV - """ - return _blob_impl.Blob.fetch(remote_path, local_path=local_path, mode="r", format="csv") - - def __call__(cls, *args, **kwargs): - """ - TODO: Is there a better way to deal with this? - - We want the behavior of Types.CSV() returns a _blob_impl.CSV, but also to be able to use - this object to wrap a _blob_impl.CSV via Types.CSV(_blob_impl.CSV()) for - serialization, type checking, etc.. - - :rtype: flytekit.common.types.impl.blobs.CSV - """ - if not args and not kwargs: - return _blob_impl.Blob.create_at_any_location(mode="w", format="csv") - else: - return super(CsvInstantiator, cls).__call__(*args, **kwargs) - - -class CSV(Blob, metaclass=CsvInstantiator): - @classmethod - def from_string(cls, string_value): - """ - :param Text string_value: - :rtype: CSV - """ - if not string_value: - _user_exceptions.FlyteValueException(string_value, "Cannot create a CSV from the provided path value.") - return cls(_blob_impl.Blob.from_string(string_value, format="csv", mode="r")) - - @classmethod - def is_castable_from(cls, other): - """ - :param flytekit.common.types.base_literal_types.FlyteSdkType other: - :rtype: bool - """ - return cls == other - - @classmethod - def from_python_std(cls, t_value): - """ - :param T t_value: It is up to each individual object as to whether or not this value can be cast. - :rtype: FlyteSdkValue - :raises: flytekit.common.exceptions.user.FlyteTypeException - """ - if t_value is None: - return _base_sdk_types.Void() - elif isinstance(t_value, _blob_impl.Blob): - if t_value.metadata.type.format != "csv": - raise _user_exceptions.FlyteValueException(t_value, "Blob is in incorrect format. Expected CSV.") - blob = t_value - else: - blob = _blob_impl.Blob.from_python_std(t_value, format="csv", mode="w") - return cls(blob) - - @classmethod - def to_flyte_literal_type(cls): - """ - :rtype: flytekit.models.types.LiteralType - """ - return _idl_types.LiteralType( - blob=_core_types.BlobType( - format="csv", - dimensionality=_core_types.BlobType.BlobDimensionality.SINGLE, - ) - ) - - @classmethod - def promote_from_model(cls, literal_model): - """ - Creates an object of this type from the model primitive defining it. - :param flytekit.models.literals.Literal literal_model: - :rtype: CSV - """ - return cls(_blob_impl.Blob.promote_from_model(literal_model.scalar.blob, mode="r")) - - @classmethod - def short_class_string(cls): - """ - :rtype: Text - """ - return "CSV" - - def __init__(self, value): - """ - :param flytekit.common.types.impl.blobs.Blob value: CSV blob value to wrap - """ - super(CSV, self).__init__(value) - - -class MultiPartCsvInstantiator(MultiPartBlobInstantiator): - @staticmethod - def create_at_known_location(location): - """ - :param Text location: - :rtype: flytekit.common.types.impl.blobs.MultiPartBlob - """ - return _blob_impl.MultiPartBlob.create_at_known_location(location, mode="w", format="csv") - - @staticmethod - def fetch(remote_path, local_path=None): - """ - :param Text remote_path: - :param Text local_path: [Optional] If specified, the MultiPartCSV is copied to this location. If specified, - this location is NOT managed and the blob will not be cleaned up upon exit. - :rtype: flytekit.common.types.impl.blobs.MultiPartCSV - """ - return _blob_impl.MultiPartBlob.fetch(remote_path, local_path=local_path, mode="r", format="csv") - - def __call__(cls, *args, **kwargs): - """ - TODO: Is there a better way to deal with this? - - We want the behavior of Types.MultiPartCSV() returns a _blob_impl.MultiPartCSV, but also to be able to use - this object to wrap a _blob_impl.MultiPartCSV via Types.MultiPartCSV(_blob_impl.MultiPartCSV()) for - serialization, type checking, etc.. - - :rtype: flytekit.common.types.impl.blobs.MultiPartCSV - """ - if not args and not kwargs: - return _blob_impl.MultiPartBlob.create_at_any_location(mode="w", format="csv") - else: - return super(MultiPartCsvInstantiator, cls).__call__(*args, **kwargs) - - -class MultiPartCSV(MultiPartBlob, metaclass=MultiPartCsvInstantiator): - @classmethod - def from_string(cls, string_value): - """ - :param Text string_value: - :rtype: MultiPartCSV - """ - if not string_value: - _user_exceptions.FlyteValueException( - string_value, - "Cannot create a MultiPartCSV from the provided path value.", - ) - return cls(_blob_impl.MultiPartBlob.from_string(string_value, format="csv", mode="r")) - - @classmethod - def is_castable_from(cls, other): - """ - :param flytekit.common.types.base_literal_types.FlyteSdkType other: - :rtype: bool - """ - return cls == other - - @classmethod - def from_python_std(cls, t_value): - """ - :param T t_value: It is up to each individual object as to whether or not this value can be cast. - :rtype: FlyteSdkValue - :raises: flytekit.common.exceptions.user.FlyteTypeException - """ - if t_value is None: - return _base_sdk_types.Void() - elif isinstance(t_value, _blob_impl.MultiPartBlob): - if t_value.metadata.type.format != "csv": - raise _user_exceptions.FlyteValueException( - t_value, "Multi Part Blob is in incorrect format. Expected CSV." - ) - blob = t_value - else: - blob = _blob_impl.MultiPartBlob.from_python_std(t_value, format="csv", mode="w") - return cls(blob) - - @classmethod - def to_flyte_literal_type(cls): - """ - :rtype: flytekit.models.types.LiteralType - """ - return _idl_types.LiteralType( - blob=_core_types.BlobType( - format="csv", - dimensionality=_core_types.BlobType.BlobDimensionality.MULTIPART, - ) - ) - - @classmethod - def promote_from_model(cls, literal_model): - """ - Creates an object of this type from the model primitive defining it. - :param flytekit.models.literals.Literal literal_model: - :rtype: MultiPartCSV - """ - return cls(_blob_impl.MultiPartBlob.promote_from_model(literal_model.scalar.blob, mode="r")) - - @classmethod - def short_class_string(cls): - """ - :rtype: Text - """ - return "MultiPartCSV" - - def __init__(self, value): - """ - :param flytekit.common.types.impl.blobs.MultiPartBlob value: MultiPartBlob value to wrap - """ - super(MultiPartCSV, self).__init__(value) diff --git a/flytekit/common/types/containers.py b/flytekit/common/types/containers.py deleted file mode 100644 index 9267273b1e..0000000000 --- a/flytekit/common/types/containers.py +++ /dev/null @@ -1,157 +0,0 @@ -import json as _json - -import six as _six - -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.types import base_sdk_types as _base_sdk_types -from flytekit.models import literals as _literals -from flytekit.models import types as _idl_types - - -class CollectionType(_base_sdk_types.FlyteSdkType): - pass - - -class TypedCollectionType(CollectionType): - @property - def sub_type(cls): - """ - :rtype: flytekit.common.types.base_sdk_types.FlyteSdkType - """ - return cls._sub_type - - def __eq__(cls, other): - return hasattr(other, "sub_type") and cls.sub_type == other.sub_type - - def __hash__(cls): - # Python 3 checks complain if hash isn't implemented at the same time as equals - return super(TypedCollectionType, cls).__hash__() - - -def List(sdk_type): - """ - :param flytekit.common.types.base_sdk_types.FlyteSdkType sdk_type: - :rtype: flytekit.common.types.base_sdk_types.FlyteSdkType - """ - - class TList(TypedListImpl): - _sub_type = sdk_type - - # TODO: Figure out generics and type-hinting - return TList - - -class ListImpl(_base_sdk_types.FlyteSdkValue, metaclass=CollectionType): - def __len__(self): - return len(self.collection.literals) - - -class TypedListImpl(ListImpl, metaclass=TypedCollectionType): - @classmethod - def from_string(cls, string_value): - """ - Load the list from a JSON formatted string. - :param Text string_value: - :rtype: ListImpl - """ - try: - items = _json.loads(string_value) - except ValueError: - raise _user_exceptions.FlyteTypeException( - _six.text_type, - cls, - additional_msg="String not parseable to json {}".format(string_value), - ) - - if type(items) != list: - raise _user_exceptions.FlyteTypeException( - _six.text_type, - cls, - additional_msg="String is not a list {}".format(string_value), - ) - - # Instead of recursively calling from_string(), we're changing to from_python_std() instead because json - # loading naturally interprets all layers, not just the outer layer. - return cls([cls.sub_type.from_python_std(i) for i in items]) - - @classmethod - def is_castable_from(cls, other): - """ - :param flytekit.common.types.base_literal_types.FlyteSdkType other: - :rtype: bool - """ - if not isinstance(type(other), TypedListImpl): - return False - return cls.sub_type.is_castable_from(other.sub_type) - - @classmethod - def from_python_std(cls, t_value): - """ - :param T t_value: It is up to each individual object as to whether or not this value can be cast. - :rtype: FlyteSdkValue - :raises: flytekit.common.exceptions.user.FlyteTypeException - """ - if t_value is None: - return _base_sdk_types.Void() - if not isinstance(t_value, list): - raise _user_exceptions.FlyteTypeException(type(t_value), list, t_value) - return cls([cls.sub_type.from_python_std(v) for v in t_value]) - - @classmethod - def to_flyte_literal_type(cls): - """ - :rtype: flytekit.models.types.LiteralType - """ - return _idl_types.LiteralType(collection_type=cls.sub_type.to_flyte_literal_type()) - - @classmethod - def promote_from_model(cls, literal_model): - """ - Creates an object of this type from the model primitive defining it. - :param flytekit.models.literals.Literal literal_model: - :rtype: TypedListImpl - """ - return cls([cls.sub_type.from_flyte_idl(l.to_flyte_idl()) for l in literal_model.collection.literals]) - - @classmethod - def short_class_string(cls): - """ - :rtype: Text - """ - return "List<{}>".format(cls.sub_type.short_class_string()) - - def __init__(self, value): - """ - :param list[flytekit.common.types.base_sdk_types.FlyteSdkValue] value: List value to wrap - """ - super(TypedListImpl, self).__init__(collection=_literals.LiteralCollection(literals=value)) - - def to_python_std(self): - """ - :rtype: list[T] - """ - return [type(self).sub_type.from_flyte_idl(l.to_flyte_idl()).to_python_std() for l in self.collection.literals] - - def short_string(self): - """ - :rtype: Text - """ - num_to_print = 5 - to_print = [v.short_string() for v in self.collection.literals[:num_to_print]] - if len(self.collection.literals) > num_to_print: - to_print.append("...") - return "{}(len={}, [{}])".format( - type(self).short_class_string(), - len(self.collection.literals), - ", ".join(to_print), - ) - - def verbose_string(self): - """ - :rtype: Text - """ - return "{}(\n\tlen={},\n\t[\n\t\t{}\n\t]\n)".format( - type(self).short_class_string(), - len(self.collection.literals), - ",\n\t\t".join("\n\t\t".join(v.verbose_string().splitlines()) for v in self.collection.literals), - ) diff --git a/flytekit/common/types/helpers.py b/flytekit/common/types/helpers.py deleted file mode 100644 index 92294f38fd..0000000000 --- a/flytekit/common/types/helpers.py +++ /dev/null @@ -1,124 +0,0 @@ -import importlib as _importlib - -import six as _six - -from flytekit.common.exceptions import scopes as _exception_scopes -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.configuration import sdk as _sdk_config -from flytekit.models import literals as _literal_models - - -class _TypeEngineLoader(object): - _LOADED_ENGINES = None - _LAST_LOADED = None - - @classmethod - def _load_engines(cls): - config = _sdk_config.TYPE_ENGINES.get() - if cls._LOADED_ENGINES is None or config != cls._LAST_LOADED: - cls._LAST_LOADED = config - cls._LOADED_ENGINES = [] - for fqdn in config: - split = fqdn.split(".") - module_path, attr = ".".join(split[:-1]), split[-1] - module = _exception_scopes.user_entry_point(_importlib.import_module)(module_path) - - if not hasattr(module, attr): - raise _user_exceptions.FlyteValueException( - module, - "Failed to load the type engine because the attribute named '{}' could not be found" - "in the module '{}'.".format(attr, module_path), - ) - - engine_impl = getattr(module, attr)() - cls._LOADED_ENGINES.append(engine_impl) - from flytekit.type_engines.default.flyte import FlyteDefaultTypeEngine as _DefaultEngine - - cls._LOADED_ENGINES.append(_DefaultEngine()) - - @classmethod - def iterate_engines_in_order(cls): - """ - :rtype: Generator[flytekit.type_engines.common.TypeEngine] - """ - cls._load_engines() - return iter(cls._LOADED_ENGINES) - - -def python_std_to_sdk_type(t): - """ - :param T t: User input. Should be of the form: Types.Integer, [Types.Integer], {Types.String: Types.Integer}, etc. - :rtype: flytekit.common.types.base_sdk_types.FlyteSdkType - """ - for e in _TypeEngineLoader.iterate_engines_in_order(): - out = e.python_std_to_sdk_type(t) - if out is not None: - return out - raise _user_exceptions.FlyteValueException(t, "Could not resolve to an SDK type for this value.") - - -def get_sdk_type_from_literal_type(literal_type): - """ - :param flytekit.models.types.LiteralType literal_type: - :rtype: flytekit.common.types.base_sdk_types.FlyteSdkType - """ - for e in _TypeEngineLoader.iterate_engines_in_order(): - out = e.get_sdk_type_from_literal_type(literal_type) - if out is not None: - return out - raise _user_exceptions.FlyteValueException( - literal_type, "Could not resolve to a type implementation for this " "value." - ) - - -def infer_sdk_type_from_literal(literal): - """ - :param flytekit.models.literals.Literal literal: - :rtype: flytekit.common.types.base_sdk_types.FlyteSdkType - """ - for e in _TypeEngineLoader.iterate_engines_in_order(): - out = e.infer_sdk_type_from_literal(literal) - if out is not None: - return out - raise _user_exceptions.FlyteValueException(literal, "Could not resolve to a type implementation for this value.") - - -def get_sdk_value_from_literal(literal, sdk_type=None): - """ - :param flytekit.models.literals.Literal literal: - :param flytekit.models.types.LiteralType sdk_type: - :rtype: flytekit.common.types.base_sdk_types.FlyteSdkValue - """ - # The spec states everything must be nullable, so if we receive a null value, swap to the null type behavior. - if sdk_type is None: - sdk_type = infer_sdk_type_from_literal(literal) - return sdk_type.from_flyte_idl(literal.to_flyte_idl()) - - -def unpack_literal_map_to_sdk_object(literal_map, type_map=None): - """ - :param lytekit.models.literals.LiteralMap literal_map: - :param dict[Text, flytekit.common.types.base_sdk_types.FlyteSdkType] type_map: Type map directing unpacking. - :rtype: dict[Text, T] - """ - type_map = type_map or {} - return {k: get_sdk_value_from_literal(v, sdk_type=type_map.get(k, None)) for k, v in literal_map.literals.items()} - - -def unpack_literal_map_to_sdk_python_std(literal_map, type_map=None): - """ - :param flytekit.models.literals.LiteralMap literal_map: Literal map containing values for unpacking. - :param dict[Text, flytekit.common.types.base_sdk_types.FlyteSdkType] type_map: Type map directing unpacking. - :rtype: dict[Text, T] - """ - return {k: v.to_python_std() for k, v in unpack_literal_map_to_sdk_object(literal_map, type_map=type_map).items()} - - -def pack_python_std_map_to_literal_map(std_map, type_map): - """ - :param dict[Text, T] std_map: - :param dict[Text, flytekit.common.types.base_sdk_types.FlyteSdkType] type_map: - :rtype: flytekit.models.literals.LiteralMap - :raises: flytekit.common.exceptions.user.FlyteTypeException - """ - return _literal_models.LiteralMap(literals={k: v.from_python_std(std_map[k]) for k, v in _six.iteritems(type_map)}) diff --git a/flytekit/common/types/impl/__init__.py b/flytekit/common/types/impl/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/flytekit/common/types/impl/blobs.py b/flytekit/common/types/impl/blobs.py deleted file mode 100644 index 45210da62a..0000000000 --- a/flytekit/common/types/impl/blobs.py +++ /dev/null @@ -1,497 +0,0 @@ -import os as _os -import shutil as _shutil -import sys as _sys -import uuid as _uuid - -import six as _six - -from flytekit.common import sdk_bases as _sdk_bases -from flytekit.common import utils as _utils -from flytekit.common.exceptions import scopes as _exception_scopes -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.interfaces.data import data_proxy as _data_proxy -from flytekit.models import literals as _literal_models -from flytekit.models.core import types as _core_types - - -class Blob(_literal_models.Blob, metaclass=_sdk_bases.ExtendedSdkType): - def __init__(self, remote_path, mode="rb", format=None): - """ - :param Text remote_path: Path to location where the Blob should be synced to. - :param Text mode: File access mode. 'a' and '+' are forbidden. A blob can only be written or read at a time. - :param Text format: Format - """ - if "+" in mode or "a" in mode or ("w" in mode and "r" in mode): - raise _user_exceptions.FlyteAssertion("A blob cannot be read and written at the same time") - self._mode = mode - self._local_path = None - self._file = None - super(Blob, self).__init__( - _literal_models.BlobMetadata( - type=_core_types.BlobType(format or "", _core_types.BlobType.BlobDimensionality.SINGLE) - ), - remote_path, - ) - - @classmethod - @_exception_scopes.system_entry_point - def from_python_std(cls, t_value, mode="wb", format=None): - """ - :param T t_value: - :param Text mode: File access mode. 'a' and '+' are forbidden. A blob can only be written or read at a time. - :param Text format: - :rtype: Blob - """ - if isinstance(t_value, (_six.text_type, str)): - if _os.path.isfile(t_value): - blob = cls.create_at_any_location(mode=mode, format=format) - blob._local_path = t_value - blob.upload() - else: - blob = cls.create_at_known_location(t_value, mode=mode, format=format) - return blob - elif isinstance(t_value, cls): - return t_value - else: - raise _user_exceptions.FlyteTypeException( - type(t_value), - {_six.text_type, str, Blob}, - received_value=t_value, - additional_msg="Unable to create Blob from user-provided value.", - ) - - @classmethod - @_exception_scopes.system_entry_point - def from_string(cls, t_value, mode="wb", format=None): - """ - :param T t_value: - :param Text mode: Read or write mode of the object. - :param Text format: - :rtype: Blob - """ - return cls.create_at_known_location(t_value, mode=mode, format=format) - - @classmethod - @_exception_scopes.system_entry_point - def create_at_known_location(cls, known_remote_location, mode="wb", format=None): - """ - :param Text known_remote_location: The location to which to write the object. Usually an s3 path. - :param Text mode: - :param Text format: - :rtype: Blob - """ - return cls(known_remote_location, mode=mode, format=format) - - @classmethod - @_exception_scopes.system_entry_point - def create_at_any_location(cls, mode="wb", format=None): - """ - :param Text mode: - :param Text format: - :rtype: Blob - """ - return cls.create_at_known_location(_data_proxy.Data.get_remote_path(), mode=mode, format=format) - - @classmethod - @_exception_scopes.system_entry_point - def fetch(cls, remote_path, local_path=None, overwrite=False, mode="rb", format=None): - """ - :param Text remote_path: The location from which to fetch the object. Usually an s3 path. - :param Text local_path: [Optional] A local path to which to download the object. If specified, the object - will not be managed and might not be cleaned up by the system upon exiting the context. - :param bool overwrite: If True, objects will be overwritten at the provided local_path in order to fetch this - object. Default is False. - :param Text mode: Read or write mode of the object. - :param Text format: Format the object is in. - :rtype: Blob - """ - blob = cls(remote_path, mode=mode, format=format) - blob.download(local_path=local_path, overwrite=overwrite) - return blob - - @classmethod - def promote_from_model(cls, model, mode="rb"): - """ - :param flytekit.models.literals.Blob model: - :param Text mode: Read or write mode of the object. - :rtype: Blob - """ - return cls(model.uri, format=model.metadata.type.format, mode=mode) - - @property - def local_path(self): - """ - Local filesystem path where the file was downloaded - :rtype: Text - """ - return self._local_path - - @property - def remote_location(self): - """ - Path to where this Blob will be synced. - :rtype: Text - """ - return self.uri - - @property - def mode(self): - """ - The mode string the Blob is associated with. - :rtype: Text - """ - return self._mode - - @_exception_scopes.system_entry_point - def __enter__(self): - """ - :rtype: typing.BinaryIO - """ - if self._file is not None: - raise _user_exceptions.FlyteAssertion("Only one reference can be open to a blob at a time.") - - if self.local_path is None: - if "r" in self.mode: - self.download() - elif "w" in self.mode: - self._generate_local_path() - - self._file = open(self.local_path, self.mode) - return self._file - - @_exception_scopes.system_entry_point - def __exit__(self, exc_type, exc_val, exc_tb): - if self._file is not None and not self._file.closed: - self._file.close() - self._file = None - if "w" in self.mode: - self.upload() - return False - - def _generate_local_path(self): - if _data_proxy.LocalWorkingDirectoryContext.get() is None: - raise _user_exceptions.FlyteAssertion( - "No temporary file system is present. Either call this method from within the " - "context of a task or surround with a 'with LocalTestFileSystem():' block. Or " - "specify a path when calling this function. Note: Cleanup is not automatic when a " - "path is specified." - ) - self._local_path = _data_proxy.LocalWorkingDirectoryContext.get().get_named_tempfile(_uuid.uuid4().hex) - - @_exception_scopes.system_entry_point - def download(self, local_path=None, overwrite=False): - """ - Alternate method, rather than the context manager interface to download the binary file to the local disk. - :param Text local_path: [Optional] If provided, the blob will be downloaded to this path. This will make the - resulting file object unmanaged and it will not be cleaned up by the system upon exiting the context. - :param bool overwrite: If true and local_path is specified, we will download the blob and - overwrite an existing file at that location. Default is False. - """ - if "r" not in self._mode: - raise _user_exceptions.FlyteAssertion("Cannot download a write-only blob!") - - if local_path: - self._local_path = local_path - - if not self.local_path: - self._generate_local_path() - - if overwrite or not _os.path.exists(self.local_path): - # TODO: Introduce system logging - # logging.info("Getting {} -> {}".format(self.remote_location, self.local_path)) - _data_proxy.Data.get_data(self.remote_location, self.local_path, is_multipart=False) - else: - raise _user_exceptions.FlyteAssertion( - "Cannot download blob to a location that already exists when overwrite is not set to True. " - "Attempted download from {} -> {}".format(self.remote_location, self.local_path) - ) - - @_exception_scopes.system_entry_point - def upload(self): - """ - Upload the blob to the remote location - """ - if "w" not in self.mode: - raise _user_exceptions.FlyteAssertion("Cannot upload a read-only blob!") - - elif not self.local_path: - raise _user_exceptions.FlyteAssertion( - "The Blob is not currently backed by a local file and therefore " - "cannot be uploaded. Please write to this Blob before attempting " - "an upload." - ) - else: - # TODO: Introduce system logging - # logging.info("Putting {} -> {}".format(self.local_path, self.remote_location)) - _data_proxy.Data.put_data(self.local_path, self.remote_location, is_multipart=False) - - -class MultiPartBlob(_literal_models.Blob, metaclass=_sdk_bases.ExtendedSdkType): - def __init__(self, remote_path, mode="rb", format=None): - """ - :param Text remote_path: Path to location where the Blob should be synced to. - :param Text mode: File access mode. 'a' and '+' are forbidden. A blob can only be written or read at a time. - :param Text format: Format of underlying blob pieces. - """ - remote_path = remote_path.strip().rstrip("/") + "/" - super(MultiPartBlob, self).__init__( - _literal_models.BlobMetadata( - type=_core_types.BlobType(format or "", _core_types.BlobType.BlobDimensionality.MULTIPART) - ), - remote_path, - ) - self._is_managed = False - self._blobs = [] - self._directory = None - self._mode = mode - - @classmethod - def promote_from_model(cls, model, mode="rb"): - """ - :param flytekit.models.literals.Blob model: - :param Text mode: File access mode. 'a' and '+' are forbidden. A blob can only be written or read at a time. - :rtype: Blob - """ - return cls(model.uri, format=model.metadata.type.format, mode=mode) - - @classmethod - @_exception_scopes.system_entry_point - def create_at_known_location(cls, known_remote_location, mode="wb", format=None): - """ - :param Text known_remote_location: The location to which to write the object. Usually an s3 path. - :param Text mode: - :param Text format: - :rtype: MultiPartBlob - """ - return cls(known_remote_location, mode=mode, format=format) - - @classmethod - @_exception_scopes.system_entry_point - def create_at_any_location(cls, mode="wb", format=None): - """ - :param Text mode: - :param Text format: - :rtype: MultiPartBlob - """ - return cls.create_at_known_location(_data_proxy.Data.get_remote_path(), mode=mode, format=format) - - @classmethod - @_exception_scopes.system_entry_point - def fetch(cls, remote_path, local_path=None, overwrite=False, mode="rb", format=None): - """ - :param Text remote_path: The location from which to fetch the object. Usually an s3 path. - :param Text local_path: [Optional] A local path to which to download the object. If specified, the object - will not be managed and might not be cleaned up by the system upon exiting the context. - :param bool overwrite: If True, objects will be overwritten at the provided local_path in order to fetch this - object. Default is False. - :param Text mode: Read or write mode of the object. - :param Text format: Format the object is in. - :rtype: MultiPartBlob - """ - blob = cls(remote_path, mode=mode, format=format) - blob.download(local_path=local_path, overwrite=overwrite) - return blob - - @classmethod - @_exception_scopes.system_entry_point - def from_python_std(cls, t_value, mode="wb", format=None): - """ - :param T t_value: - :param Text mode: Read or write mode of the object. - :param Text format: - :rtype: MultiPartBlob - """ - if isinstance(t_value, (str, _six.text_type)): - if _os.path.isdir(t_value): - # TODO: Infer format - blob = cls.create_at_any_location(mode=mode, format=format) - blob._directory = _utils.Directory(t_value) - blob.upload() - else: - blob = cls.create_at_known_location(t_value, mode=mode, format=format) - return blob - elif isinstance(t_value, cls): - return t_value - else: - raise _user_exceptions.FlyteTypeException( - type(t_value), - {str, _six.text_type, MultiPartBlob}, - received_value=t_value, - additional_msg="Unable to create Blob from user-provided value.", - ) - - @classmethod - @_exception_scopes.system_entry_point - def from_string(cls, t_value, mode="wb", format=None): - """ - :param T t_value: - :param Text mode: Read or write mode of the object. - :param Text format: - :rtype: MultiPartBlob - """ - return cls.create_at_known_location(t_value, mode=mode, format=format) - - @_exception_scopes.system_entry_point - def __enter__(self): - """ - :rtype: list[typing.BinaryIO] - """ - if "r" not in self.mode: - raise _user_exceptions.FlyteAssertion("Do not enter context to write to directory. Call create_piece") - - try: - if not self._directory: - if _data_proxy.LocalWorkingDirectoryContext.get() is None: - raise _user_exceptions.FlyteAssertion( - "No temporary file system is present. Either call this method from within the " - "context of a task or surround with a 'with LocalTestFileSystem():' block. Or " - "specify a path when calling this function. Note: Cleanup is not automatic when a " - "path is specified." - ) - self._directory = _utils.AutoDeletingTempDir( - _uuid.uuid4().hex, - tmp_dir=_data_proxy.LocalWorkingDirectoryContext.get().name, - ) - self._is_managed = True - self._directory.__enter__() - # TODO: Introduce system logging - # logging.info("Copying recursively {} -> {}".format(self.remote_location, self.local_path)) - _data_proxy.Data.get_data(self.remote_location, self.local_path, is_multipart=True) - - # Read the files into blobs in case-insensitive lexicographically ascending orders - self._blobs = [] - file_handles = [] - for local_path in sorted(self._directory.list_dir(), key=lambda x: x.lower()): - b = Blob( - _os.path.join(self.remote_location, _os.path.basename(local_path)), - mode=self.mode, - ) - b._local_path = local_path - file_handles.append(b.__enter__()) - self._blobs.append(b) - - return file_handles - except Exception: - # Exit is idempotent so close partially opened context that way - exc_type, exc_obj, exc_tb = _sys.exc_info() - self.__exit__(exc_type, exc_obj, exc_tb) - raise - - @_exception_scopes.system_entry_point - def __exit__(self, exc_type, exc_val, exc_tb): - for blob in self._blobs: - blob.__exit__(exc_type, exc_val, exc_tb) - self._blobs = [] - if self._is_managed: - self._directory.__exit__(exc_type, exc_val, exc_tb) - self._directory = None - self._is_managed = False - return False - - @property - def local_path(self): - """ - Local filesystem path where the file was downloaded - :rtype: Text - """ - if not self._directory: - return None - return self._directory.name - - @property - def remote_location(self): - """ - Path to where this MultiPartBlob will be synced. - :rtype: Text - """ - return self.uri - - @property - def mode(self): - """ - The mode string the MultiPartBlob is associated with. - :rtype: Text - """ - return self._mode - - @_exception_scopes.system_entry_point - def create_part(self, name=None): - """ - Method which will return a Blob object for writing into a multi-part blob. - - :param Text name: [optional] If name is provided, it is a specific partition name to place as part of the - multipart blob. When we read blobs from a multipart, it will be read in lexicographic order so this can be - used to enforce ordering. If not provided, the name is randomly generated. - :rtype: Blob - """ - if "w" not in self.mode: - raise _user_exceptions.FlyteAssertion("Cannot create a blob in a read-only multipart blob") - if name is None: - name = _uuid.uuid4().hex - if ":" in name or "/" in name: - raise _user_exceptions.FlyteAssertion( - name, - "Cannot create a part of a multi-part object with ':' or '/' in the name.", - ) - return Blob.create_at_known_location( - _os.path.join(self.remote_location, name), - mode=self.mode, - format=self.metadata.type.format, - ) - - @_exception_scopes.system_entry_point - def download(self, local_path=None, overwrite=False): - """ - Forces the download of the remote multi-part blob to the local machine. - :param Text local_path: [Optional] If provided, the blob pieces will be downloaded to this path. This will - make the resulting file objects unmanaged and it will not be cleaned up by the system upon exiting the - context. - :param bool overwrite: If true and local_path is specified, we will download the blob pieces and - overwrite any existing files at that location. Default is False. - """ - if "r" not in self.mode: - raise _user_exceptions.FlyteAssertion("Cannot download a write-only object!") - - if local_path: - self._is_managed = False - elif _data_proxy.LocalWorkingDirectoryContext.get() is None: - raise _user_exceptions.FlyteAssertion( - "No temporary file system is present. Either call this method from within the " - "context of a task or surround with a 'with LocalTestFileSystem():' block. Or " - "specify a path when calling this function. Note: Cleanup is not automatic when a " - "path is specified." - ) - else: - local_path = _data_proxy.LocalWorkingDirectoryContext.get().get_named_tempfile(_uuid.uuid4().hex) - - path_exists = _os.path.exists(local_path) - if not path_exists or overwrite: - if path_exists: - _shutil.rmtree(local_path) - _os.makedirs(local_path) - self._directory = _utils.Directory(local_path) - _data_proxy.Data.get_data(self.remote_location, self.local_path, is_multipart=True) - else: - raise _user_exceptions.FlyteAssertion( - "Cannot download multi-part blob to a location that already exists when overwrite is not set to True. " - "Attempted download from {} -> {}".format(self.remote_location, self.local_path) - ) - - @_exception_scopes.system_entry_point - def upload(self): - """ - Upload the multi-part blob to the remote location - """ - if "w" not in self.mode: - raise _user_exceptions.FlyteAssertion("Cannot upload a read-only multi-part blob!") - - elif not self.local_path: - raise _user_exceptions.FlyteAssertion( - "The multi-part blob is not currently backed by a local directoru " - "and therefore cannot be uploaded. Please write to this before " - "attempting an upload." - ) - else: - # TODO: Introduce system logging - # logging.info("Putting {} -> {}".format(self.local_path, self.remote_location)) - _data_proxy.Data.put_data(self.local_path, self.remote_location, is_multipart=True) diff --git a/flytekit/common/types/impl/schema.py b/flytekit/common/types/impl/schema.py deleted file mode 100644 index 04e4109d1c..0000000000 --- a/flytekit/common/types/impl/schema.py +++ /dev/null @@ -1,995 +0,0 @@ -import collections as _collections -import os as _os -import uuid as _uuid - -import six as _six - -from flytekit.common import sdk_bases as _sdk_bases -from flytekit.common import utils as _utils -from flytekit.common.exceptions import scopes as _exception_scopes -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.types import base_sdk_types as _base_sdk_types -from flytekit.common.types import helpers as _helpers -from flytekit.common.types import primitives as _primitives -from flytekit.common.types.impl import blobs as _blob_impl -from flytekit.configuration import sdk as _sdk_config -from flytekit.interfaces.data import data_proxy as _data_proxy -from flytekit.models import literals as _literal_models -from flytekit.models import types as _type_models -from flytekit.plugins import numpy as _np -from flytekit.plugins import pandas as _pd - -# Note: For now, this is only for basic type-checking. We need not differentiate between TINYINT, BIGINT, -# and INT or DOUBLE and FLOAT, VARCHAR and STRING, etc. as we will unpack into appropriate Python -# objects anyway. If we work on managed tables, these more specific type specifications might become necessary. -_SUPPORTED_LITERAL_TYPE_TO_PANDAS_TYPES = None - - -def get_supported_literal_types_to_pandas_types(): - global _SUPPORTED_LITERAL_TYPE_TO_PANDAS_TYPES - if _SUPPORTED_LITERAL_TYPE_TO_PANDAS_TYPES is None: - _SUPPORTED_LITERAL_TYPE_TO_PANDAS_TYPES = { - _primitives.Integer.to_flyte_literal_type(): {_np.int32, _np.int64, _np.uint32, _np.uint64}, - _primitives.Float.to_flyte_literal_type(): {_np.float32, _np.float64}, - _primitives.Boolean.to_flyte_literal_type(): {_np.bool}, - _primitives.Datetime.to_flyte_literal_type(): {_np.datetime64}, - _primitives.Timedelta.to_flyte_literal_type(): {_np.timedelta64}, - _primitives.String.to_flyte_literal_type(): {_np.object_, _np.str_, _np.string_}, - } - return _SUPPORTED_LITERAL_TYPE_TO_PANDAS_TYPES - - -_ALLOWED_PARTITION_TYPES = {str, int} - -# Hive currently has limitations where column headers are not stored when writing to an overwrite directory. There is -# an open proposal (https://issues.apache.org/jira/browse/HIVE-12860) to improve this. Until then, we have this -# work-around where we create an external table with the appropriate schema and write the data to our desired -# location. The issue here is that the table information in the meta-store might not get cleaned up during a partial -# failure. -_HIVE_QUERY_FORMATTER = """ - {stage_query_str} - - CREATE TEMPORARY TABLE {table}_tmp AS {query_str}; - CREATE EXTERNAL TABLE {table} LIKE {table}_tmp STORED AS PARQUET; - ALTER TABLE {table} SET LOCATION '{url}'; - - INSERT OVERWRITE TABLE {table} - SELECT - {columnar_query} - FROM {table}_tmp; - DROP TABLE {table}; - """ - -# Once https://issues.apache.org/jira/browse/HIVE-12860 is resolved. We will prefer the following syntax because it -# guarantees cleanup on partial failures. -_HIVE_QUERY_FORMATTER_V2 = """ - CREATE TEMPORARY TABLE {table} AS {query_str}; - - INSERT OVERWRITE DIRECTORY '{url}' STORED AS PARQUET - SELECT {columnar_query} - FROM {table}; - """ - -# Set location in both parts of this query so in case of a partial failure, we will always have some data backing a -# partition. -_WRITE_HIVE_PARTITION_QUERY_FORMATTER = """ - ALTER TABLE {write_table} ADD IF NOT EXISTS {partition_string} LOCATION '{url}'; - - ALTER TABLE {write_table} {partition_string} SET LOCATION '{url}'; - """ - - -def _format_insert_partition_query(table_name, partition_string, remote_location): - table_pieces = table_name.split(".") - if len(table_pieces) > 1: - # Hive shell commands don't allow us to alter tables and select databases in the table specification. So - # we split the table name and use the 'use' command to choose the correct database. - prefix = "use {};\n".format(table_pieces[0]) - table_name = ".".join(table_pieces[1:]) - else: - prefix = "" - - return prefix + _WRITE_HIVE_PARTITION_QUERY_FORMATTER.format( - write_table=table_name, partition_string=partition_string, url=remote_location - ) - - -class _SchemaIO(object): - def __init__(self, schema_instance, local_dir, mode): - """ - :param Schema schema_instance: - :param flytekit.common.utils.Directory local_dir: - :param Text mode: - """ - self._schema = schema_instance - self._local_dir = local_dir - self._chunks = [] - self._index = 0 - self._mode = mode - - def _access_guard(self): - if not self._schema: - raise _user_exceptions.FlyteAssertion( - "Schema IO object has already been closed. Cannot access chunk_count property." - ) - - @_exception_scopes.system_entry_point - def iter_chunks(self, *args, **kwargs): - raise _user_exceptions.FlyteAssertion("{} is write only.".format(self._schema)) - - @_exception_scopes.system_entry_point - def read(self, *args, **kwargs): - raise _user_exceptions.FlyteAssertion("{} is write only.".format(self._schema)) - - @_exception_scopes.system_entry_point - def write(self, *args, **kwargs): - raise _user_exceptions.FlyteAssertion("{} is read only.".format(self._schema)) - - @_exception_scopes.system_entry_point - def close(self): - self._schema = None - self._local_dir = None - self._chunks = None - self._index = 0 - - @property - @_exception_scopes.system_entry_point - def chunk_count(self): - self._access_guard() - return len(self._chunks) - - @_exception_scopes.system_entry_point - def seek(self, index): - self._access_guard() - if index < 0 or index > self.chunk_count: - raise _user_exceptions.FlyteValueException( - index, - "Attempting to seek to a chunk that is out of range. Allowed range is [0, {}]".format(self.chunk_count), - ) - self._index = index - - @_exception_scopes.system_entry_point - def tell(self): - return self._index - - def __repr__(self): - return "{mode} IO Object for {type} @ {location}".format( - type=self._schema.type, location=self._schema.remote_prefix, mode=self._mode - ) - - -class _SchemaReader(_SchemaIO): - def __init__(self, schema_instance, local_dir): - """ - :param Schema schema_instance: - :param flytekit.common.utils.Directory local_dir: - """ - super(_SchemaReader, self).__init__(schema_instance, local_dir, "Read-Only") - self.reset_chunks() - - @_exception_scopes.system_entry_point - def reset_chunks(self): - self._chunks = sorted(self._local_dir.list_dir()) - - @_exception_scopes.system_entry_point - def iter_chunks(self, columns=None, **kwargs): - self._access_guard() - while self._index < len(self._chunks): - chunk = self.read(columns=columns, concat=False, **kwargs) - if chunk is not None: - yield chunk - - @staticmethod - def _read_parquet_with_type_promotion_override(chunk, columns, parquet_engine): - """ - This wrapper function of pd.read_parquet() is a hack intended to fix the type promotion problem - when using fastparquet as the underlying parquet engine. - - When using fastparquet, boolean columns containing None values will be promoted to float16 columns. - This becomes problematic when users want to write the dataframe back into parquet - file because float16 (halffloat) is not a supported type in parquet spec. In this function, we detect - such columns and do override the type promotion. - """ - df = None - - if parquet_engine == "fastparquet": - import fastparquet.thrift_structures as _ts - from fastparquet import ParquetFile as _ParquetFile - - # https://github.com/dask/fastparquet/issues/414#issuecomment-478983811 - df = _pd.read_parquet(chunk, columns=columns, engine=parquet_engine, index=False) - df_column_types = df.dtypes - pf = _ParquetFile(chunk) - schema_column_dtypes = {l.name: l.type for l in list(pf.schema.schema_elements)} - - for idx in df_column_types[df_column_types == "float16"].index.tolist(): - # A hacky way to get the string representations of the column types of a parquet schema - # Reference: - # https://github.com/dask/fastparquet/blob/f4ecc67f50e7bf98b2d0099c9589c615ea4b06aa/fastparquet/schema.py - if _ts.parquet_thrift.Type._VALUES_TO_NAMES[schema_column_dtypes[idx]] == "BOOLEAN": - df[idx] = df[idx].astype("object") - df[idx].replace({0: False, 1: True, _pd.np.nan: None}, inplace=True) - - else: - df = _pd.read_parquet(chunk, columns=columns, engine=parquet_engine) - - return df - - @_exception_scopes.system_entry_point - def read(self, columns=None, concat=False, truncate_extra_columns=True, **kwargs): - """ - When this function is called, one chunk will be read and received as a Pandas data frame. Once all chunks - have been read, this function will return None. - - :param list[Text] columns: A list of columns to read. They must be a subset of the columns - defined for the Schema object. If specified, truncate_extra_columns must be True. - :param bool concat: If true, the entire object will be returned in one large data frame. - :param bool truncate_extra_columns: If true, only columns from the underlying parquet file will be read if - they are specified as columns in the schema object (except for empty schemas which will read all columns - regardless). If false, if there are additional columns in the underlying parquet file, they will also be - read. - :rtype: pandas.DataFrame - """ - if columns is not None and truncate_extra_columns is False: - raise _user_exceptions.FlyteAssertion( - "When reading a schema object, it is not possible to both specify a set of columns to read and " - "additionally not truncate_extra_columns. Either columns must not be specified or " - "truncate_extra_columns must be set to True (or not specified)." - ) - - self._access_guard() - - parquet_engine = _sdk_config.PARQUET_ENGINE.get() - if parquet_engine not in {"fastparquet", "pyarrow"}: - raise _user_exceptions.FlyteAssertion( - "environment variable parquet_engine must be one of 'pyarrow', 'fastparquet', or be unset" - ) - - df_out = None - if not columns: - columns = list(self._schema.type.sdk_columns.keys()) - - if len(columns) == 0 or truncate_extra_columns is False: - columns = None - - if concat: - frames = [ - # A hacky hack - # TODO: follow up the issue opened in the fastparquet repo for a more general fix - # issue URL: - _SchemaReader._read_parquet_with_type_promotion_override( - chunk=chunk, columns=columns, parquet_engine=parquet_engine - ) - # _pd.read_parquet(chunk, columns=columns, engine=parquet_engine) - for chunk in self._chunks[self._index :] - if _os.path.getsize(chunk) > 0 - ] - if len(frames) == 1: - df_out = frames[0] - elif len(frames) > 1: - df_out = _pd.concat(frames, copy=True) - self._index = len(self._chunks) - else: - while self._index < len(self._chunks) and df_out is None: - # Skip empty chunks so the user appears to have a continuous stream of data. - if _os.path.getsize(self._chunks[self._index]) > 0: - df_out = _SchemaReader._read_parquet_with_type_promotion_override( - chunk=self._chunks[self._index], columns=columns, parquet_engine=parquet_engine, **kwargs - ) - self._index += 1 - - if df_out is not None: - self._schema.compare_dataframe_to_schema(df_out, read=True, column_subset=columns) - - # Make sure the columns are renamed to exactly what the user specifies. This prevents unexpected - # unicode v. string mismatches. Also, if a schema is mapped with strict_names=False, the input might - # have totally different names. - user_columns = columns or _six.iterkeys(self._schema.type.sdk_columns) - # User-specified columns may or may not be unicode - # Since, in python 2, dictionary does a transparent translation between unicode and str for the key, - # (https://stackoverflow.com/a/24532329) - # we use this characteristic to create a trivial lookup dictionary, to make sure we can use either - # unicode or str to lookup, but get back whatever type the user actually used - user_column_dict = {c: c for c in user_columns} - if len(self._schema.type.columns) > 0: - # Avoid using pandas.DataFrame.rename() as this function incurs significant memory overhead - df_out.columns = [ - user_column_dict[col] if col in user_columns else col for col in df_out.columns.values - ] - return df_out - - -class _SchemaWriter(_SchemaIO): - def __init__(self, schema_instance, local_dir): - """ - :param Schema schema_instance: - :param flytekit.common.utils.Directory local_dir: - :param Text mode: - """ - super(_SchemaWriter, self).__init__(schema_instance, local_dir, "Write-Only") - - @_exception_scopes.system_entry_point - def close(self): - """ - Closes the writing IO context and uploads data to s3. - """ - try: - # TODO: Introduce system logging - # logging.info("Copying recursively {} -> {}".format(self._local_dir.name, self._schema.remote_prefix)) - _data_proxy.Data.put_data(self._local_dir.name, self._schema.remote_prefix, is_multipart=True) - finally: - super(_SchemaWriter, self).close() - - @_exception_scopes.system_entry_point - def write(self, data_frame, coerce_timestamps="us", allow_truncated_timestamps=False): - """ - Writes data frame as a chunk to the local directory owned by the Schema object. Will later be uploaded to s3. - - :param pandas.DataFrame data_frame: data frame to write as parquet - :param Text coerce_timestamps: format to store timestamp in parquet. 'us', 'ms', 's' are allowed values. - Note: if your timestamps will lose data due to the coercion, your write will fail! Nanoseconds are - problematic in the Parquet format and will not work. See allow_truncated_timestamps. - :param bool allow_truncated_timestamps: default False. Allow truncation when coercing timestamps to a coarser - resolution. - """ - self._access_guard() - if not isinstance(data_frame, _pd.DataFrame): - raise _user_exceptions.FlyteTypeException( - expected_type=_pd.DataFrame, - received_type=type(data_frame), - received_value=data_frame, - additional_msg="Only pandas DataFrame objects can be written to a Schema object", - ) - - self._schema.compare_dataframe_to_schema(data_frame) - all_columns = list(data_frame.columns.values) - - # Convert all columns to unicode as pyarrow's parquet reader can not handle mixed strings and unicode. - # Since columns from Hive are returned as unicode, if a user wants to add a column to a dataframe returned from - # Hive, then output the new data, the user would have to provide a unicode column name which is unnatural. - unicode_columns = [_six.text_type(col) for col in all_columns] - data_frame.columns = unicode_columns - try: - filename = self._local_dir.get_named_tempfile(_os.path.join(str(self._index).zfill(6))) - data_frame.to_parquet( - filename, - coerce_timestamps=coerce_timestamps, - allow_truncated_timestamps=allow_truncated_timestamps, - ) - if self._index == len(self._chunks): - self._chunks.append(filename) - self._index += 1 - finally: - # Return to old names to prevent odd behavior with user. - data_frame.columns = unicode_columns - - -class _SchemaBackingMpBlob(_blob_impl.MultiPartBlob): - @property - def directory(self): - """ - :rtype: flytekit.common.utils.Directory - """ - return self._directory - - def __enter__(self): - if not self.local_path: - if _data_proxy.LocalWorkingDirectoryContext.get() is None: - raise _user_exceptions.FlyteAssertion( - "No temporary file system is present. Either call this method from within the " - "context of a task or surround with a 'with LocalTestFileSystem():' block. Or " - "specify a path when calling this function." - ) - self._directory = _utils.AutoDeletingTempDir( - _uuid.uuid4().hex, - tmp_dir=_data_proxy.LocalWorkingDirectoryContext.get().name, - ) - self._is_managed = True - self._directory.__enter__() - - if "r" in self.mode: - _data_proxy.Data.get_data(self.remote_location, self.local_path, is_multipart=True) - - def __exit__(self, exc_type, exc_val, exc_tb): - if "w" in self.mode: - _data_proxy.Data.put_data(self.local_path, self.remote_location, is_multipart=True) - return super(_SchemaBackingMpBlob, self).__exit__(exc_type, exc_val, exc_tb) - - -class SchemaType(_type_models.SchemaType, metaclass=_sdk_bases.ExtendedSdkType): - _LITERAL_TYPE_TO_PROTO_ENUM = { - _primitives.Integer.to_flyte_literal_type(): _type_models.SchemaType.SchemaColumn.SchemaColumnType.INTEGER, - _primitives.Float.to_flyte_literal_type(): _type_models.SchemaType.SchemaColumn.SchemaColumnType.FLOAT, - _primitives.Boolean.to_flyte_literal_type(): _type_models.SchemaType.SchemaColumn.SchemaColumnType.BOOLEAN, - _primitives.Datetime.to_flyte_literal_type(): _type_models.SchemaType.SchemaColumn.SchemaColumnType.DATETIME, - _primitives.Timedelta.to_flyte_literal_type(): _type_models.SchemaType.SchemaColumn.SchemaColumnType.DURATION, - _primitives.String.to_flyte_literal_type(): _type_models.SchemaType.SchemaColumn.SchemaColumnType.STRING, - } - - def __init__(self, columns=None): - super(SchemaType, self).__init__(None) - self._set_columns(columns or []) - - @property - def sdk_columns(self): - """ - This is an ordered dictionary so iterating over it will be in the order columns were specified in the - constructor. - :rtype: dict[Text, flytekit.common.types.base_sdk_types.FlyteSdkType] - """ - return self._sdk_columns - - @property - def columns(self): - """ - :rtype: list[flytekit.models.types.SchemaType.SchemaColumn] - """ - return [ - _type_models.SchemaType.SchemaColumn(n, type(self)._LITERAL_TYPE_TO_PROTO_ENUM[v.to_flyte_literal_type()]) - for n, v in _six.iteritems(self.sdk_columns) - ] - - @classmethod - def promote_from_model(cls, model): - """ - :param flytekit.models.types.SchemaType model: - :rtype: SchemaType - """ - _PROTO_ENUM_TO_SDK_TYPE = { - _type_models.SchemaType.SchemaColumn.SchemaColumnType.INTEGER: _helpers.get_sdk_type_from_literal_type( - _primitives.Integer.to_flyte_literal_type() - ), - _type_models.SchemaType.SchemaColumn.SchemaColumnType.FLOAT: _helpers.get_sdk_type_from_literal_type( - _primitives.Float.to_flyte_literal_type() - ), - _type_models.SchemaType.SchemaColumn.SchemaColumnType.BOOLEAN: _helpers.get_sdk_type_from_literal_type( - _primitives.Boolean.to_flyte_literal_type() - ), - _type_models.SchemaType.SchemaColumn.SchemaColumnType.DATETIME: _helpers.get_sdk_type_from_literal_type( - _primitives.Datetime.to_flyte_literal_type() - ), - _type_models.SchemaType.SchemaColumn.SchemaColumnType.DURATION: _helpers.get_sdk_type_from_literal_type( - _primitives.Timedelta.to_flyte_literal_type() - ), - _type_models.SchemaType.SchemaColumn.SchemaColumnType.STRING: _helpers.get_sdk_type_from_literal_type( - _primitives.String.to_flyte_literal_type() - ), - } - return cls([(c.name, _PROTO_ENUM_TO_SDK_TYPE[c.type]) for c in model.columns]) - - def _set_columns(self, columns): - names_seen = set() - for column in columns: - if not isinstance(column, tuple): - raise _user_exceptions.FlyteValueException( - column, - "When specifying a Schema type with a known set of columns. Each column must be " - "specified as a tuple in the form ('name', type).", - ) - if len(column) != 2: - raise _user_exceptions.FlyteValueException( - column, - "When specifying a Schema type with a known set of columns. Each column must be " - "specified as a tuple in the form ('name', type).", - ) - name, sdk_type = column - sdk_type = _helpers.python_std_to_sdk_type(sdk_type) - - if not isinstance(name, (str, _six.text_type)): - additional_msg = ( - "When specifying a Schema type with a known set of columns, the first element in" - " each tuple must be text." - ) - raise _user_exceptions.FlyteTypeException( - received_type=type(name), - received_value=name, - expected_type={str, _six.text_type}, - additional_msg=additional_msg, - ) - - if ( - not isinstance(sdk_type, _base_sdk_types.FlyteSdkType) - or sdk_type.to_flyte_literal_type() not in get_supported_literal_types_to_pandas_types() - ): - additional_msg = ( - "When specifying a Schema type with a known set of columns, the second element of " - "each tuple must be a supported type. Failed for column: {name}".format(name=name) - ) - raise _user_exceptions.FlyteTypeException( - expected_type=list(get_supported_literal_types_to_pandas_types().keys()), - received_type=sdk_type, - additional_msg=additional_msg, - ) - - if name in names_seen: - raise ValueError( - "The column name {name} was specified multiple times when instantiating the " - "Schema.".format(name=name) - ) - names_seen.add(name) - - self._sdk_columns = _collections.OrderedDict(columns) - - -class Schema(_literal_models.Schema, metaclass=_sdk_bases.ExtendedSdkType): - def __init__(self, remote_path, mode="rb", schema_type=None): - """ - :param Text remote_path: - :param Text mode: - :param SchemaType schema_type: [Optional] If specified, the schema will be forced to conform to this type. If - not specified, the schema will be considered generic. - """ - self._mp_blob = _SchemaBackingMpBlob(remote_path, mode=mode) - super(Schema, self).__init__(self._mp_blob.uri, schema_type or SchemaType()) - self._io_object = None - - @classmethod - def promote_from_model(cls, model): - """ - :param flytekit.models.literals.Schema model: - :rtype: Schema - """ - return cls(model.uri, schema_type=SchemaType.promote_from_model(model.type)) - - @classmethod - @_exception_scopes.system_entry_point - def create_at_known_location(cls, known_remote_location, mode="wb", schema_type=None): - """ - :param Text known_remote_location: The location to which to write the object. Usually an s3 path. - :param Text mode: - :param SchemaType schema_type: [Optional] If specified, the schema will be forced to conform to this type. If - not specified, the schema will be considered generic. - :rtype: Schema - """ - return cls(known_remote_location, mode=mode, schema_type=schema_type) - - @classmethod - @_exception_scopes.system_entry_point - def create_at_any_location(cls, mode="wb", schema_type=None): - """ - :param Text mode: - :param SchemaType schema_type: [Optional] If specified, the schema will be forced to conform to this type. If - not specified, the schema will be considered generic. - :rtype: Schema - """ - return cls.create_at_known_location(_data_proxy.Data.get_remote_path(), mode=mode, schema_type=schema_type) - - @classmethod - @_exception_scopes.system_entry_point - def fetch(cls, remote_path, local_path=None, overwrite=False, mode="rb", schema_type=None): - """ - :param Text remote_path: The location from which to fetch the object. Usually an s3 path. - :param Text local_path: [Optional] A local path to which to download the object. If specified, the object - will not be managed and might not be cleaned up by the system upon exiting the context. - :param bool overwrite: If True, objects will be overwritten at the provided local_path in order to fetch this - object. Default is False. - :param Text mode: Read or write mode of the object. - :param SchemaType schema_type: [Optional] If specified, the schema will be forced to conform to this type. If - not specified, the schema will be considered generic. - :rtype: Schema - """ - schema = cls(remote_path, mode=mode, schema_type=schema_type) - schema.download(local_path=local_path, overwrite=overwrite) - return schema - - @classmethod - @_exception_scopes.system_entry_point - def from_python_std(cls, t_value, schema_type=None): - """ - :param T t_value: - :param SchemaType schema_type: [Optional] If specified, we will ensure - :rtype: Schema - """ - if isinstance(t_value, (str, _six.text_type)): - if _os.path.isdir(t_value): - schema = cls.create_at_any_location(schema_type=schema_type) - schema.multipart_blob._directory = _utils.Directory(t_value) - schema.upload() - else: - schema = cls.create_at_known_location(t_value, schema_type=schema_type) - return schema - elif isinstance(t_value, cls): - return t_value - elif isinstance(t_value, _pd.DataFrame): - # Accepts a pandas dataframe and converts to a Schema object - o = cls.create_at_any_location(schema_type=schema_type) - with o as w: - w.write(t_value) - return o - elif isinstance(t_value, list): - # Accepts a list of pandas dataframe and converts to a Schema object - o = cls.create_at_any_location(schema_type=schema_type) - with o as w: - for x in t_value: - if isinstance(x, _pd.DataFrame): - w.write(x) - else: - raise _user_exceptions.FlyteTypeException( - type(t_value), - {str, _six.text_type, Schema}, - received_value=x, - additional_msg="A Schema object can only be create from a pandas DataFrame or a list of pandas DataFrame.", - ) - return o - else: - raise _user_exceptions.FlyteTypeException( - type(t_value), - {str, _six.text_type, Schema}, - received_value=t_value, - additional_msg="Unable to create Schema from user-provided value.", - ) - - @classmethod - def from_string(cls, string_value, schema_type=None): - """ - :param Text string_value: - :param SchemaType schema_type: - :rtype: Schema - """ - if not string_value: - _user_exceptions.FlyteValueException(string_value, "Cannot create a Schema from an empty path") - return cls.create_at_known_location(string_value, schema_type=schema_type) - - @classmethod - @_exception_scopes.system_entry_point - def create_from_hive_query( - cls, - select_query, - stage_query=None, - schema_to_table_name_map=None, - schema_type=None, - known_location=None, - ): - """ - Returns a query that can be submitted to Hive and produce the desired output. It also returns a properly-typed - schema object. - - :param Text select_query: Query for selecting data from Hive - :param Text stage_query: Query for building temporary tables on Hive. - Runs before the select query. Temporary tables are supported but CTEs are not supported. - :param Dict[Text, Text] schema_to_table_name_map: A map of column names in the schema to the column names - returned from the select query - :param Text known_location: create the schema object at a known s3 location. - :param SchemaType schema_type: [Optional] If specified, the schema will be forced to conform to this type. If - not specified, the schema will be considered generic. - :return: Schema, Text - """ - schema_object = cls( - known_location or _data_proxy.Data.get_remote_directory(), - mode="wb", - schema_type=schema_type, - ) - - if len(schema_object.type.sdk_columns) > 0: - identity_dict = {n: n for n in _six.iterkeys(schema_object.type.sdk_columns)} - identity_dict.update(schema_to_table_name_map or {}) - schema_to_table_name_map = identity_dict - - columnar_clauses = [] - for name, sdk_type in _six.iteritems(schema_object.type.sdk_columns): - if sdk_type == _primitives.Float: - columnar_clauses.append( - "CAST({table_column_name} as double) {schema_name}".format( - table_column_name=schema_to_table_name_map[name], - schema_name=name, - ) - ) - else: - columnar_clauses.append( - "{table_column_name} as {schema_name}".format( - table_column_name=schema_to_table_name_map[name], - schema_name=name, - ) - ) - columnar_query = ",\n\t\t".join(columnar_clauses) - else: - columnar_query = "*" - - stage_query_str = _six.text_type(stage_query or "") - # the stage query should always end with a semicolon - stage_query_str = stage_query_str if stage_query_str.endswith(";") else (stage_query_str + ";") - query = _HIVE_QUERY_FORMATTER.format( - url=schema_object.remote_location, - stage_query_str=stage_query_str, - query_str=select_query.strip().strip(";"), - columnar_query=columnar_query, - table=_uuid.uuid4().hex, - ) - return schema_object, query - - @property - def local_path(self): - """ - Local filesystem path where the file was downloaded - :rtype: Text - """ - return self._mp_blob.local_path - - @property - def remote_location(self): - """ - Path to where this MultiPartBlob will be synced. This is needed for reverse compatibility. - :rtype: Text - """ - return self.uri - - @property - def remote_prefix(self): - """ - Path to where this MultiPartBlob will be synced. This is needed for reverse compatibility. - :rtype: Text - """ - return self.uri - - @property - def uri(self): - """ - Path to where this MultiPartBlob will be synced. - :rtype: Text - """ - return self.multipart_blob.uri - - @property - def mode(self): - """ - The mode string the MultiPartBlob is associated with. - :rtype: Text - """ - return self._mp_blob.mode - - @property - def type(self): - """ - The schema type definition associated with this object. - :rtype: SchemaType - """ - return self._type - - @property - def multipart_blob(self): - """ - :rtype: flytekit.common.types.impl.blobs.MultiPartBlob - """ - return self._mp_blob - - @_exception_scopes.system_entry_point - def __enter__(self): - """ - :rtype: _SchemaIO - """ - if self._io_object is not None: - raise _user_exceptions.FlyteAssertion( - "The context of a schema can only be entered once at a time. Make sure the previous " - "'with' block has been exited." - ) - - self._mp_blob.__enter__() - if "r" in self.mode: - self._io_object = _SchemaReader(self, self.multipart_blob.directory) - else: - self._io_object = _SchemaWriter(self, self.multipart_blob.directory) - return self._io_object - - @_exception_scopes.system_entry_point - def __exit__(self, exc_type, exc_val, exc_tb): - self._io_object = None - return self._mp_blob.__exit__(exc_type, exc_val, exc_tb) - - def __repr__(self): - return "Schema({columns}) @ {location} ({mode})".format( - columns=self.type.columns, - location=self.remote_prefix, - mode="read-only" if "r" in self.mode else "write-only", - ) - - @_exception_scopes.system_entry_point - def download(self, local_path=None, overwrite=False): - """ - :param Text local_path: [Optional] A local path to which to download the object. If specified, the object - will not be managed and might not be cleaned up by the system upon exiting the context. - :param bool overwrite: If True, objects will be overwritten at the provided local_path in order to fetch this - object. Default is False. - :rtype: Schema - """ - self.multipart_blob.download(local_path=local_path, overwrite=overwrite) - - @_exception_scopes.system_entry_point - def get_write_partition_to_hive_table_query( - self, - table_name, - partitions=None, - schema_to_table_name_map=None, - partitions_in_table=False, - append_to_partition=False, - ): - """ - Returns a Hive query string that will update the metatable to point to the data as the new partition. - - :param Text table_name: - :param dict[Text, T] partitions: A dictionary mapping table partition key names to the values matching this - partition. - :param dict[Text, Text] schema_to_table_name_map: Mapping of names in current schema to table in which it is - being inserted. Currently not supported. Must be None. - :param bool partitions_in_table: Whether or not the partition columns exist in the data being submitted. - Currently not supported. Must be false - :param bool append_to_partition: Whether or not to append new values to a partition. Currently not supported. - :return: Text - """ - partition_string = "" - where_string = "" - identity_dict = {n: n for n in _six.iterkeys(self.type.sdk_columns)} - identity_dict.update(schema_to_table_name_map or {}) - schema_to_table_name_map = identity_dict - table_to_schema_name_map = {v: k for k, v in _six.iteritems(schema_to_table_name_map)} - - if partitions: - partition_conditions = [] - for partition_name, partition_value in _six.iteritems(partitions): - if not isinstance(partition_name, (str, _six.text_type)): - raise _user_exceptions.FlyteTypeException( - expected_type={str, _six.text_type}, - received_type=type(partition_name), - received_value=partition_name, - additional_msg="All partition names must be type str.", - ) - if type(partition_value) not in _ALLOWED_PARTITION_TYPES: - raise _user_exceptions.FlyteTypeException( - expected_type=_ALLOWED_PARTITION_TYPES, - received_type=type(partition_value), - received_value=partition_value, - additional_msg="Partition {name} has an unsupported type.".format(name=partition_name), - ) - - # We need the string to be quoted in the query, so let's take repr of it. - if isinstance(partition_value, (str, _six.text_type)): - partition_value = repr(partition_value) - partition_conditions.append( - "{partition_name} = {partition_value}".format( - partition_name=partition_name, partition_value=partition_value - ) - ) - partition_formatter = "PARTITION (\n\t{conditions}\n)" - partition_string = partition_formatter.format(conditions=",\n\t".join(partition_conditions)) - - if partitions_in_table and partitions: - where_clauses = [] - for partition_name, partition_value in partitions: - where_clauses.append( - "\n\t\t{schema_name} = {value_str} AND ".format( - schema_name=table_to_schema_name_map[partition_name], - value_str=partition_value, - ) - ) - where_string = "WHERE\n\t\t{where_clauses}".format(where_clauses=" AND\n\t\t".join(where_clauses)) - - if where_string or partitions_in_table: - raise _user_exceptions.FlyteAssertion( - "Currently, the partition values should not be present in the schema pushed to Hive." - ) - if append_to_partition: - raise _user_exceptions.FlyteAssertion( - "Currently, partitions can only be overwritten, they cannot be appended." - ) - if not partitions: - raise _user_exceptions.FlyteAssertion( - "Currently, partition values MUST be specified for writing to a table." - ) - - return _format_insert_partition_query( - remote_location=self.remote_location, - table_name=table_name, - partition_string=partition_string, - ) - - def compare_dataframe_to_schema(self, data_frame, column_subset=None, read=False): - """ - Do necessary type checking of a pandas data frame. Raise exception if it doesn't match. - :param pandas.DateFrame data_frame: data frame to type check - :param list[Text] column_subset: - :param bool read: Used to alter error message for more clarity. - """ - all_columns = list(data_frame.columns.values) - schema_column_names = list(self.type.sdk_columns.keys()) - - # Skip checking if we have a generic schema type (no specified columns) - if not schema_column_names: - return - - # If we specify a subset of columns, ensure they all exist and then only take those columns - if column_subset is not None: - schema_column_names = [] - failed_columns = [] - for column in column_subset: - if column not in self.type.sdk_columns: - failed_columns.append(column) - else: - schema_column_names.append(column) - - if len(failed_columns) > 0: - additional_msg = "" - raise _user_exceptions.FlyteAssertion( - "{} was/where requested but could not be found in the schema: {}.{}".format( - failed_columns, self.type.sdk_columns, additional_msg - ) - ) - - if not all(c in all_columns for c in schema_column_names): - raise _user_exceptions.FlyteTypeException( - expected_type=self.type.sdk_columns, - received_type=data_frame.columns, - additional_msg="Mismatch between the data frame's column names {} and schema's column names {} " - "with strict_names=True.".format(all_columns, schema_column_names), - ) - - # This only iterates if the Schema has specified columns. - for name in schema_column_names: - literal_type = self.type.sdk_columns[name].to_flyte_literal_type() - dtype = data_frame[name].dtype - - # TODO np.issubdtype is deprecated. Replace it - if all( - not _np.issubdtype(dtype, allowed_type) - for allowed_type in get_supported_literal_types_to_pandas_types()[literal_type] - ): - if read: - read_or_write_msg = "read data frame object from schema" - else: - read_or_write_msg = "write data frame object to schema" - additional_msg = ( - "Cannot {read_write} because the types do not match. Column " - "'{name}' did not pass type checking. Note: If your " - "column contains null values, the types might not transition as expected between parquet and " - "pandas. For more information, see: " - "http://arrow.apache.org/docs/python/pandas.html#arrow-pandas-conversion".format( - read_write=read_or_write_msg, name=name - ) - ) - raise _user_exceptions.FlyteTypeException( - expected_type=get_supported_literal_types_to_pandas_types()[literal_type], - received_type=dtype, - additional_msg=additional_msg, - ) - - def cast_to(self, other_type): - """ - :param SchemaType other_type: - :rtype: Schema - """ - if len(other_type.sdk_columns) > 0: - for k, v in _six.iteritems(other_type.sdk_columns): - if k not in self.type.sdk_columns: - raise _user_exceptions.FlyteTypeException( - self.type, - other_type, - additional_msg="Cannot cast because a required column '{}' was not found.".format(k), - received_value=self, - ) - if ( - not isinstance(v, _base_sdk_types.FlyteSdkType) - or v.to_flyte_literal_type() != self.type.sdk_columns[k].to_flyte_literal_type() - ): - raise _user_exceptions.FlyteTypeException( - self.type.sdk_columns[k], - v, - additional_msg="Cannot cast because the column type for column '{}' does not match.".format(k), - ) - return Schema(self.remote_location, mode=self.mode, schema_type=other_type) - - @_exception_scopes.system_entry_point - def upload(self): - """ - Upload the schema to the remote location - """ - if "w" not in self.mode: - raise _user_exceptions.FlyteAssertion("Cannot upload a read-only schema!") - - elif not self.local_path: - raise _user_exceptions.FlyteAssertion( - "The schema is not currently backed by a local directory " - "and therefore cannot be uploaded. Please write to this before " - "attempting an upload." - ) - else: - # TODO: Introduce system logging - # logging.info("Putting {} -> {}".format(self.local_path, self.remote_location)) - _data_proxy.Data.put_data(self.local_path, self.remote_location, is_multipart=True) diff --git a/flytekit/common/types/primitives.py b/flytekit/common/types/primitives.py deleted file mode 100644 index 446140dc2c..0000000000 --- a/flytekit/common/types/primitives.py +++ /dev/null @@ -1,595 +0,0 @@ -import datetime as _datetime -import json as _json -import typing - -import six as _six -from dateutil import parser as _parser -from google.protobuf import json_format as _json_format -from google.protobuf import struct_pb2 as _struct -from pytimeparse import parse as _parse_duration_string - -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.types import base_sdk_types as _base_sdk_types -from flytekit.models import literals as _literals -from flytekit.models import types as _idl_types - - -class Integer(_base_sdk_types.FlyteSdkValue): - @classmethod - def from_string(cls, string_value): - """ - :param Text string_value: - :rtype: Integer - """ - try: - return cls(int(string_value)) - except (ValueError, TypeError): - raise _user_exceptions.FlyteTypeException( - _six.text_type, - int, - additional_msg="String not castable to Integer SDK type:" " {}".format(string_value), - ) - - @classmethod - def is_castable_from(cls, other): - """ - :param flytekit.common.types.base_literal_types.FlyteSdkType other: - :rtype: bool - """ - return cls == other - - @classmethod - def from_python_std(cls, t_value): - """ - :param T t_value: It is up to each individual object as to whether or not this value can be cast. - :rtype: FlyteSdkValue - :raises: flytekit.common.exceptions.user.FlyteTypeException - """ - if t_value is None: - return _base_sdk_types.Void() - if type(t_value) not in _six.integer_types: - raise _user_exceptions.FlyteTypeException(type(t_value), _six.integer_types, t_value) - return cls(t_value) - - @classmethod - def to_flyte_literal_type(cls): - """ - :rtype: flytekit.models.types.LiteralType - """ - return _idl_types.LiteralType(simple=_idl_types.SimpleType.INTEGER) - - @classmethod - def promote_from_model(cls, literal_model): - """ - Creates an object of this type from the model primitive defining it. - :param flytekit.models.literals.Literal literal_model: - :rtype: Integer - """ - return cls(literal_model.scalar.primitive.integer) - - @classmethod - def short_class_string(cls): - """ - :rtype: Text - """ - return "Integer" - - def __init__(self, value): - """ - :param int value: Int value to wrap - """ - super(Integer, self).__init__(scalar=_literals.Scalar(primitive=_literals.Primitive(integer=value))) - - def to_python_std(self): - """ - :rtype: int - """ - return self.scalar.primitive.integer - - def short_string(self): - """ - :rtype: Text - """ - return "Integer({})".format(self.scalar.primitive.integer) - - -class Float(_base_sdk_types.FlyteSdkValue): - @classmethod - def from_string(cls, string_value): - """ - :param Text string_value: - :rtype: Float - """ - try: - return cls(float(string_value)) - except ValueError: - raise _user_exceptions.FlyteTypeException( - _six.text_type, - float, - additional_msg="String not castable to Float SDK type:" " {}".format(string_value), - ) - - @classmethod - def is_castable_from(cls, other): - """ - :param flytekit.common.types.base_literal_types.FlyteSdkType other: - :rtype: bool - """ - return cls == other - - @classmethod - def from_python_std(cls, t_value): - """ - :param T t_value: It is up to each individual object as to whether or not this value can be cast. - :rtype: FlyteSdkValue - :raises: flytekit.common.exceptions.user.FlyteTypeException - """ - if t_value is None: - return _base_sdk_types.Void() - if type(t_value) != float: - raise _user_exceptions.FlyteTypeException(type(t_value), float, t_value) - return cls(t_value) - - @classmethod - def to_flyte_literal_type(cls): - """ - :rtype: flytekit.models.types.LiteralType - """ - return _idl_types.LiteralType(simple=_idl_types.SimpleType.FLOAT) - - @classmethod - def promote_from_model(cls, literal_model): - """ - Creates an object of this type from the model primitive defining it. - :param flytekit.models.literals.Literal literal_model: - :rtype: Float - """ - return cls(literal_model.scalar.primitive.float_value) - - @classmethod - def short_class_string(cls): - """ - :rtype: Text - """ - return "Float" - - def __init__(self, value): - """ - :param float value: value to wrap - """ - super(Float, self).__init__(scalar=_literals.Scalar(primitive=_literals.Primitive(float_value=value))) - - def to_python_std(self): - """ - :rtype: float - """ - return self.scalar.primitive.float_value - - def short_string(self): - """ - :rtype: Text - """ - return "Float({})".format(self.scalar.primitive.float_value) - - -class Boolean(_base_sdk_types.FlyteSdkValue): - @classmethod - def from_string(cls, string_value): - """ - :param Text string_value: - :rtype: Boolean - """ - if string_value == "1" or string_value.lower() == "true": - return cls(True) - elif string_value == "0" or string_value.lower() == "false": - return cls(False) - raise _user_exceptions.FlyteTypeException( - _six.text_type, - bool, - additional_msg="String not castable to Boolean SDK " "type: {}".format(string_value), - ) - - @classmethod - def is_castable_from(cls, other): - """ - :param flytekit.common.types.base_literal_types.FlyteSdkType other: - :rtype: bool - """ - return cls == other - - @classmethod - def from_python_std(cls, t_value): - """ - :param T t_value: It is up to each individual object as to whether or not this value can be cast. - :rtype: FlyteSdkValue - :raises: flytekit.common.exceptions.user.FlyteTypeException - """ - if t_value is None: - return _base_sdk_types.Void() - if type(t_value) != bool: - raise _user_exceptions.FlyteTypeException(type(t_value), bool, t_value) - return cls(t_value) - - @classmethod - def to_flyte_literal_type(cls): - """ - :rtype: flytekit.models.types.LiteralType - """ - return _idl_types.LiteralType(simple=_idl_types.SimpleType.BOOLEAN) - - @classmethod - def promote_from_model(cls, literal_model): - """ - Creates an object of this type from the model primitive defining it. - :param flytekit.models.literals.Literal literal_model: - :rtype: bool - """ - return cls(literal_model.scalar.primitive.boolean) - - @classmethod - def short_class_string(cls): - """ - :rtype: Text - """ - return "Boolean" - - def __init__(self, value): - """ - :param bool value: value to wrap - """ - super(Boolean, self).__init__(scalar=_literals.Scalar(primitive=_literals.Primitive(boolean=value))) - - def to_python_std(self): - """ - :rtype: bool - """ - return self.scalar.primitive.boolean - - def short_string(self): - """ - :rtype: Text - """ - return "Boolean({})".format(self.scalar.primitive.boolean) - - -class String(_base_sdk_types.FlyteSdkValue): - @classmethod - def from_string(cls, string_value): - """ - :param Text string_value: - :rtype: String - """ - if type(string_value) == dict or type(string_value) == list: - raise _user_exceptions.FlyteTypeException( - type(string_value), - _six.text_type, - additional_msg="Should not cast native Python type to string {}".format(string_value), - ) - return cls(string_value) - - @classmethod - def is_castable_from(cls, other): - """ - :param flytekit.common.types.base_literal_types.FlyteSdkType other: - :rtype: bool - """ - return cls == other - - @classmethod - def from_python_std(cls, t_value): - """ - Creates an object of this type from the model primitive defining it. - :param T t_value: It is up to each individual object as to whether or not this value can be cast. - :rtype: FlyteSdkValue - :raises: flytekit.common.exceptions.user.FlyteTypeException - """ - if t_value is None: - return _base_sdk_types.Void() - if type(t_value) not in set([str, _six.text_type]): - raise _user_exceptions.FlyteTypeException(type(t_value), set([str, _six.text_type]), t_value) - return cls(t_value) - - @classmethod - def to_flyte_literal_type(cls): - """ - :rtype: flytekit.models.types.LiteralType - """ - return _idl_types.LiteralType(simple=_idl_types.SimpleType.STRING) - - @classmethod - def promote_from_model(cls, literal_model): - """ - Creates an object of this type from the model primitive defining it. - :param flytekit.models.literals.Literal literal_model: - :rtype: String - """ - return cls(literal_model.scalar.primitive.string_value) - - @classmethod - def short_class_string(cls): - """ - :rtype: Text - """ - return "String" - - def __init__(self, value): - """ - :param Text value: value to wrap - """ - super(String, self).__init__(scalar=_literals.Scalar(primitive=_literals.Primitive(string_value=value))) - - def to_python_std(self): - """ - :rtype: Text - """ - return self.scalar.primitive.string_value - - def short_string(self): - """ - :rtype: Text - """ - _TRUNCATE_LENGTH = 100 - return "String('{}'{})".format( - self.scalar.primitive.string_value[:_TRUNCATE_LENGTH], - " ..." if len(self.scalar.primitive.string_value) > _TRUNCATE_LENGTH else "", - ) - - def verbose_string(self): - """ - :rtype: Text - """ - return "String('{}')".format(self.scalar.primitive.string_value) - - -class Datetime(_base_sdk_types.FlyteSdkValue): - @classmethod - def from_string(cls, string_value): - """ - :param Text string_value: - :rtype: Datetime - """ - try: - python_std_datetime = _parser.parse(string_value) - except ValueError: - raise _user_exceptions.FlyteTypeException( - _six.text_type, - _datetime.datetime, - additional_msg="String not castable to Datetime " "SDK type: {}".format(string_value), - ) - - return cls.from_python_std(python_std_datetime) - - @classmethod - def is_castable_from(cls, other): - """ - :param flytekit.common.types.base_literal_types.FlyteSdkType other: - :rtype: bool - """ - return cls == other - - @classmethod - def from_python_std(cls, t_value): - """ - :param T t_value: It is up to each individual object as to whether or not this value can be cast. - :rtype: FlyteSdkValue - :raises: flytekit.common.exceptions.user.FlyteTypeException - """ - if t_value is None: - return _base_sdk_types.Void() - elif type(t_value) != _datetime.datetime: - raise _user_exceptions.FlyteTypeException(type(t_value), _datetime.datetime, t_value) - elif t_value.tzinfo is None: - raise _user_exceptions.FlyteValueException( - t_value, - "Datetime objects in Flyte must be timezone aware. " "tzinfo was found to be None.", - ) - return cls(t_value) - - @classmethod - def to_flyte_literal_type(cls): - """ - :rtype: flytekit.models.types.LiteralType - """ - return _idl_types.LiteralType(simple=_idl_types.SimpleType.DATETIME) - - @classmethod - def promote_from_model(cls, literal_model): - """ - Creates an object of this type from the model primitive defining it. - :param flytekit.models.literals.Literal literal_model: - :rtype: Datetime - """ - return cls(literal_model.scalar.primitive.datetime) - - @classmethod - def short_class_string(cls): - """ - :rtype: Text - """ - return "Datetime" - - def __init__(self, value): - """ - :param datetime.datetime value: value to wrap - """ - super(Datetime, self).__init__(scalar=_literals.Scalar(primitive=_literals.Primitive(datetime=value))) - - def to_python_std(self): - """ - :rtype: datetime.datetime - """ - return self.scalar.primitive.datetime - - def short_string(self): - """ - :rtype: Text - """ - return "Datetime({})".format(_six.text_type(self.scalar.primitive.datetime)) - - -class Timedelta(_base_sdk_types.FlyteSdkValue): - @classmethod - def from_string(cls, string_value): - """ - :param Text string_value: Uses https://github.com/wroberts/pytimeparse for parsing - :rtype: Timedelta - """ - td = _parse_duration_string(string_value) - if td is None: - raise _user_exceptions.FlyteTypeException( - _six.text_type, - _datetime.timedelta, - additional_msg="Could not convert string to" " time delta: {}".format(string_value), - ) - return cls.from_python_std(_datetime.timedelta(seconds=td)) - - @classmethod - def is_castable_from(cls, other): - """ - :param flytekit.common.types.base_literal_types.FlyteSdkType other: - :rtype: bool - """ - return cls == other - - @classmethod - def from_python_std(cls, t_value): - """ - :param T t_value: It is up to each individual object as to whether or not this value can be cast. - :rtype: FlyteSdkValue - :raises: flytekit.common.exceptions.user.FlyteTypeException - """ - if t_value is None: - return _base_sdk_types.Void() - elif type(t_value) != _datetime.timedelta: - raise _user_exceptions.FlyteTypeException(type(t_value), _datetime.timedelta, t_value) - - return cls(t_value) - - @classmethod - def to_flyte_literal_type(cls): - """ - :rtype: flytekit.models.types.LiteralType - """ - return _idl_types.LiteralType(simple=_idl_types.SimpleType.DURATION) - - @classmethod - def promote_from_model(cls, literal_model): - """ - Creates an object of this type from the model primitive defining it. - :param flytekit.models.literals.Literal literal_model: - :rtype: Timedelta - """ - return cls(literal_model.scalar.primitive.duration) - - @classmethod - def short_class_string(cls): - """ - :rtype: Text - """ - return "Timedelta" - - def __init__(self, value): - """ - :param datetime.timedelta value: value to wrap - """ - super(Timedelta, self).__init__(scalar=_literals.Scalar(primitive=_literals.Primitive(duration=value))) - - def to_python_std(self): - """ - :rtype: datetime.timedelta - """ - return self.scalar.primitive.duration - - def short_string(self): - """ - :rtype: Text - """ - return "Timedelta({})".format(self.scalar.primitive.duration) - - -class Generic(_base_sdk_types.FlyteSdkValue): - @classmethod - def from_string(cls, string_value): - """ - :param Text string_value: Should be a JSON formatted string - :rtype: Generic - """ - try: - t = _json_format.Parse(string_value, _struct.Struct()) - except Exception: - raise _user_exceptions.FlyteValueException(string_value, "Could not be parsed from JSON.") - return cls(t) - - @classmethod - def is_castable_from(cls, other): - """ - :param flytekit.common.types.base_literal_types.FlyteSdkType other: - :rtype: bool - """ - return cls == other - - @classmethod - def from_python_std(cls, t_value): - """ - :param T t_value: It is up to each individual object as to whether or not this value can be cast. - :rtype: FlyteSdkValue - :raises: flytekit.common.exceptions.user.FlyteTypeException - """ - if t_value is None: - return _base_sdk_types.Void() - elif not isinstance(t_value, dict): - raise _user_exceptions.FlyteTypeException(type(t_value), dict, t_value) - - try: - t = _json.dumps(t_value) - except Exception: - raise _user_exceptions.FlyteValueException(t_value, "Is not JSON serializable.") - - return cls(_json_format.Parse(t, _struct.Struct())) - - @classmethod - def to_flyte_literal_type(cls, metadata: typing.Dict = None): - """ - :rtype: flytekit.models.types.LiteralType - """ - return _idl_types.LiteralType(simple=_idl_types.SimpleType.STRUCT, metadata=metadata) - - @classmethod - def promote_from_model(cls, literal_model): - """ - Creates an object of this type from the model primitive defining it. - :param flytekit.models.literals.Literal literal_model: - :rtype: Generic - """ - return cls(literal_model.scalar.generic) - - @classmethod - def short_class_string(cls): - """ - :rtype: Text - """ - return "Generic" - - def __init__(self, value): - """ - :param _struct.Struct value: value to wrap - """ - super(Generic, self).__init__(scalar=_literals.Scalar(generic=value)) - - def to_python_std(self): - """ - :rtype: dict[Text, T] - """ - return _json.loads(_json_format.MessageToJson(self.scalar.generic)) - - def short_string(self): - """ - :rtype: Text - """ - return "Generic({})".format(self.to_python_std()) - - def long_string(self): - """ - :rtype: Text - """ - return "Generic({})".format(self.to_python_std()) diff --git a/flytekit/common/types/proto.py b/flytekit/common/types/proto.py deleted file mode 100644 index 2c3423ccfb..0000000000 --- a/flytekit/common/types/proto.py +++ /dev/null @@ -1,319 +0,0 @@ -import base64 as _base64 -from typing import Type, Union - -import six as _six -from google.protobuf import reflection as _proto_reflection -from google.protobuf.json_format import Error -from google.protobuf.json_format import MessageToDict as _MessageToDict -from google.protobuf.json_format import ParseDict as _ParseDict -from google.protobuf.reflection import GeneratedProtocolMessageType -from google.protobuf.struct_pb2 import Struct - -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.types import base_sdk_types as _base_sdk_types -from flytekit.models import literals as _literals -from flytekit.models import types as _idl_types -from flytekit.models.common import FlyteIdlEntity, FlyteType -from flytekit.models.types import LiteralType - -ProtobufT = Type[_proto_reflection.GeneratedProtocolMessageType] - - -class ProtobufType(_base_sdk_types.FlyteSdkType): - _pb_type = Struct - - @property - def pb_type(cls) -> GeneratedProtocolMessageType: - """ - :rtype: GeneratedProtocolMessageType - """ - return cls._pb_type - - @property - def descriptor(cls): - """ - :rtype: Text - """ - return "{}.{}".format(cls.pb_type.__module__, cls.pb_type.__name__) - - @property - def tag(cls): - """ - :rtype: Text - """ - return "{}{}".format(Protobuf.TAG_PREFIX, cls.descriptor) - - -class Protobuf(_base_sdk_types.FlyteSdkValue, metaclass=ProtobufType): - PB_FIELD_KEY = "pb_type" - TAG_PREFIX = "{}=".format(PB_FIELD_KEY) - - def __init__(self, pb_object: Union[GeneratedProtocolMessageType, FlyteIdlEntity]): - """ - :param Union[T, FlyteIdlEntity] pb_object: - """ - v = pb_object - # This section converts an existing proto object (or a subclass of) to the right type expected by this instance - # of GenericProto. GenericProto can be used with any protobuf type (not restricted to FlyteType). This makes it - # a bit tricky to figure out the right version of the underlying raw proto class to use to populate the final - # struct. - # If the provided object has to_flyte_idl(), call it to produce a raw proto. - if isinstance(pb_object, FlyteIdlEntity): - v = pb_object.to_flyte_idl() - - # A check to ensure the raw proto (v) is of the correct expected type. This also performs one final attempt to - # convert it to the correct type by leveraging from_flyte_idl (implemented by all FlyteTypes) in case this class - # is initialized with one. - expected_type = type(self).pb_type - if expected_type != type(v) and expected_type != type(pb_object): - if isinstance(type(self).pb_type, FlyteType): - v = expected_type.from_flyte_idl(v).to_flyte_idl() - else: - raise _user_exceptions.FlyteTypeException( - received_type=type(pb_object), expected_type=expected_type, received_value=pb_object - ) - data = v.SerializeToString() - super(Protobuf, self).__init__( - scalar=_literals.Scalar( - binary=_literals.Binary(value=bytes(data) if _six.PY2 else data, tag=type(self).tag) - ) - ) - - @classmethod - def from_string(cls, string_value): - """ - :param Text string_value: b64 encoded string of bytes - :rtype: Protobuf - """ - try: - decoded = _base64.b64decode(string_value) - except TypeError: - raise _user_exceptions.FlyteValueException(string_value, "The string is not valid base64-encoded.") - pb_obj = cls.pb_type() - pb_obj.ParseFromString(decoded) - return cls(pb_obj) - - @classmethod - def is_castable_from(cls, other): - """ - :param flytekit.common.types.base_literal_types.FlyteSdkType other: - :rtype: bool - """ - return isinstance(other, ProtobufType) and other.pb_type is cls.pb_type - - @classmethod - def from_python_std(cls, t_value): - """ - :param T t_value: It is up to each individual object as to whether or not this value can be cast. - :rtype: _base_sdk_types.FlyteSdkValue - :raises: flytekit.common.exceptions.user.FlyteTypeException - """ - if t_value is None: - return _base_sdk_types.Void() - elif isinstance(t_value, cls.pb_type) or isinstance(t_value, FlyteIdlEntity): - return cls(t_value) - else: - raise _user_exceptions.FlyteTypeException(type(t_value), cls.pb_type, received_value=t_value) - - @classmethod - def to_flyte_literal_type(cls): - """ - :rtype: flytekit.models.types.LiteralType - """ - return _idl_types.LiteralType(simple=_idl_types.SimpleType.BINARY, metadata={cls.PB_FIELD_KEY: cls.descriptor}) - - @classmethod - def promote_from_model(cls, literal_model): - """ - Creates an object of this type from the model primitive defining it. - :param flytekit.models.literals.Literal literal_model: - :rtype: Protobuf - """ - if literal_model.scalar.binary.tag != cls.tag: - raise _user_exceptions.FlyteTypeException( - literal_model.scalar.binary.tag, - cls.pb_type, - received_value=_base64.b64encode(literal_model.scalar.binary.value), - additional_msg="Can not deserialize as proto tags don't match.", - ) - pb_obj = cls.pb_type() - pb_obj.ParseFromString(literal_model.scalar.binary.value) - return cls(pb_obj) - - @classmethod - def short_class_string(cls): - """ - :rtype: Text - """ - return "Types.Proto({})".format(cls.descriptor) - - def to_python_std(self): - """ - :returns: The protobuf object as defined by the user. - :rtype: T - """ - pb_obj = type(self).pb_type() - pb_obj.ParseFromString(self.scalar.binary.value) - return pb_obj - - def short_string(self): - """ - :rtype: Text - """ - return "{}".format(self.to_python_std()) - - -def create_protobuf(pb_type: Type[GeneratedProtocolMessageType]) -> Type[Protobuf]: - """ - :param Type[GeneratedProtocolMessageType] pb_type: - :rtype: Type[Protobuf] - """ - if not isinstance(pb_type, _proto_reflection.GeneratedProtocolMessageType): - raise _user_exceptions.FlyteTypeException( - expected_type=_proto_reflection.GeneratedProtocolMessageType, - received_type=type(pb_type), - received_value=pb_type, - ) - - class _Protobuf(Protobuf): - _pb_type = pb_type - - return _Protobuf - - -class GenericProtobuf(_base_sdk_types.FlyteSdkValue, metaclass=ProtobufType): - PB_FIELD_KEY = "pb_type" - TAG_PREFIX = "{}=".format(PB_FIELD_KEY) - - def __init__(self, pb_object: Union[GeneratedProtocolMessageType, FlyteIdlEntity]): - """ - :param Union[T, FlyteIdlEntity] pb_object: - """ - struct = Struct() - v = pb_object - - # This section converts an existing proto object (or a subclass of) to the right type expected by this instance - # of GenericProto. GenericProto can be used with any protobuf type (not restricted to FlyteType). This makes it - # a bit tricky to figure out the right version of the underlying raw proto class to use to populate the final - # struct. - # If the provided object has to_flyte_idl(), call it to produce a raw proto. - if isinstance(pb_object, FlyteIdlEntity): - v = pb_object.to_flyte_idl() - - # A check to ensure the raw proto (v) is of the correct expected type. This also performs one final attempt to - # convert it to the correct type by leveraging from_flyte_idl (implemented by all FlyteTypes) in case this class - # is initialized with one. - expected_type = type(self).pb_type - if expected_type != type(v) and expected_type != type(pb_object): - if isinstance(type(self).pb_type, FlyteType): - v = expected_type.from_flyte_idl(v).to_flyte_idl() - else: - raise _user_exceptions.FlyteTypeException( - received_type=type(pb_object), expected_type=expected_type, received_value=pb_object - ) - - struct.update(_MessageToDict(v)) - super().__init__(scalar=_literals.Scalar(generic=struct)) - - @classmethod - def is_castable_from(cls, other): - """ - :param flytekit.common.types.base_literal_types.FlyteSdkType other: - :rtype: bool - """ - return isinstance(other, ProtobufType) and other.pb_type is cls.pb_type - - @classmethod - def from_python_std(cls, t_value: Union[GeneratedProtocolMessageType, FlyteIdlEntity]): - """ - :param Union[T, FlyteIdlEntity] t_value: It is up to each individual object as to whether or not this value can be cast. - :rtype: _base_sdk_types.FlyteSdkValue - :raises: flytekit.common.exceptions.user.FlyteTypeException - """ - if t_value is None: - return _base_sdk_types.Void() - elif isinstance(t_value, cls.pb_type) or isinstance(t_value, FlyteIdlEntity): - return cls(t_value) - else: - raise _user_exceptions.FlyteTypeException(type(t_value), cls.pb_type, received_value=t_value) - - @classmethod - def to_flyte_literal_type(cls) -> LiteralType: - """ - :rtype: flytekit.models.types.LiteralType - """ - return _idl_types.LiteralType(simple=_idl_types.SimpleType.STRUCT, metadata={cls.PB_FIELD_KEY: cls.descriptor}) - - @classmethod - def promote_from_model(cls, literal_model): - """ - Creates an object of this type from the model primitive defining it. - :param flytekit.models.literals.Literal literal_model: - :rtype: Protobuf - """ - pb_obj = cls.pb_type() - try: - dictionary = _MessageToDict(literal_model.scalar.generic) - pb_obj = _ParseDict(dictionary, pb_obj) - except Error as err: - raise _user_exceptions.FlyteTypeException( - received_type="generic", - expected_type=cls.pb_type, - received_value=_base64.b64encode(literal_model.scalar.generic), - additional_msg=f"Can not deserialize. Error: {err.__str__()}", - ) - - return cls(pb_obj) - - @classmethod - def short_class_string(cls) -> str: - """ - :rtype: Text - """ - return "Types.GenericProto({})".format(cls.descriptor) - - def to_python_std(self): - """ - :returns: The protobuf object as defined by the user. - :rtype: T - """ - pb_obj = type(self).pb_type() - try: - dictionary = _MessageToDict(self.scalar.generic) - pb_obj = _ParseDict(dictionary, pb_obj) - except Error as err: - raise _user_exceptions.FlyteTypeException( - received_type="generic", - expected_type=type(self).pb_type, - received_value=_base64.b64encode(self.scalar.generic), - additional_msg=f"Can not deserialize. Error: {err.__str__()}", - ) - return pb_obj - - def short_string(self) -> str: - """ - :rtype: Text - """ - return "{}".format(self.to_python_std()) - - -def create_generic(pb_type: Type[GeneratedProtocolMessageType]) -> Type[GenericProtobuf]: - """ - Creates a generic protobuf type that represents protobuf type ProtobufT and that will get serialized into a struct. - - :param Type[GeneratedProtocolMessageType] pb_type: - :rtype: Type[GenericProtobuf] - """ - if not isinstance(pb_type, _proto_reflection.GeneratedProtocolMessageType) and not issubclass( - pb_type, FlyteIdlEntity - ): - raise _user_exceptions.FlyteTypeException( - expected_type=_proto_reflection.GeneratedProtocolMessageType, - received_type=type(pb_type), - received_value=pb_type, - ) - - class _Protobuf(GenericProtobuf): - _pb_type = pb_type - - return _Protobuf diff --git a/flytekit/common/types/schema.py b/flytekit/common/types/schema.py deleted file mode 100644 index eaf38d1c88..0000000000 --- a/flytekit/common/types/schema.py +++ /dev/null @@ -1,189 +0,0 @@ -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.types import base_sdk_types as _base_sdk_types -from flytekit.common.types.impl import schema as _schema_impl -from flytekit.models import literals as _literals -from flytekit.models import types as _idl_types - - -class SchemaInstantiator(_base_sdk_types.InstantiableType): - def create_at_known_location(cls, location): - """ - :param Text location: - :rtype: flytekit.common.types.impl.schema.Schema - """ - return _schema_impl.Schema.create_at_known_location(location, mode="wb", schema_type=cls.schema_type) - - def fetch(cls, remote_path, local_path=None): - """ - :param Text remote_path: - :param Text local_path: [Optional] If specified, the Schema is copied to this location. If specified, - this location is NOT managed and the schema will not be cleaned up upon exit. - :rtype: flytekit.common.types.impl.schema.Schema - """ - return _schema_impl.Schema.fetch(remote_path, mode="rb", local_path=local_path, schema_type=cls.schema_type) - - def create(cls): - """ - :rtype: flytekit.common.types.impl.schema.Schema - """ - return _schema_impl.Schema.create_at_any_location(mode="wb", schema_type=cls.schema_type) - - def create_from_hive_query( - cls, - select_query, - stage_query=None, - schema_to_table_name_map=None, - known_location=None, - ): - """ - Returns a query that can be submitted to Hive and produce the desired output. It also returns a properly-typed - schema object. - - :param Text select_query: Query for selecting data from Hive - :param Text stage_query: Query for building temporary tables on Hive. - Runs before the select query. Temporary tables are supported but CTEs are not supported. - :param Dict[Text, Text] schema_to_table_name_map: A map of column names in the schema to the column names - returned from the select query - :param Text known_location: create the schema object at a known s3 location. - :rtype: flytekit.common.types.impl.schema.Schema, Text - """ - return _schema_impl.Schema.create_from_hive_query( - select_query=select_query, - stage_query=stage_query, - schema_to_table_name_map=schema_to_table_name_map, - known_location=known_location, - schema_type=cls.schema_type, - ) - - def __call__(cls, *args, **kwargs): - """ - TODO: Is there a better way to deal with this? - - :rtype: flytekit.common.types.impl.schema.Schema - """ - if not args and not kwargs: - return _schema_impl.Schema.create_at_any_location(mode="wb", schema_type=cls.schema_type) - else: - return super(SchemaInstantiator, cls).__call__(*args, **kwargs) - - @property - def schema_type(cls): - """ - :rtype: _schema_impl.SchemaType - """ - return cls._schema_type - - @property - def columns(cls): - """ - :rtype: dict[Text, flytekit.common.types.base_sdk_types.FlyteSdkType] - """ - return cls.schema_type.sdk_columns - - -class Schema(_base_sdk_types.FlyteSdkValue, metaclass=SchemaInstantiator): - @classmethod - def from_string(cls, string_value): - """ - :param Text string_value: - :rtype: Schema - """ - if not string_value: - _user_exceptions.FlyteValueException(string_value, "Cannot create a Schema from an empty path") - return cls(_schema_impl.Schema.from_string(string_value, schema_type=cls.schema_type)) - - @classmethod - def is_castable_from(cls, other): - """ - :param flytekit.common.types.base_literal_types.FlyteSdkType other: - :rtype: bool - """ - return cls == other - - @classmethod - def from_python_std(cls, t_value): - """ - :param T t_value: It is up to each individual object as to whether or not this value can be cast. - :rtype: FlyteSdkValue - :raises: flytekit.common.exceptions.user.FlyteTypeException - """ - if t_value is None: - return _base_sdk_types.Void() - elif isinstance(t_value, _schema_impl.Schema): - schema = t_value.cast_to(cls.schema_type) - else: - schema = _schema_impl.Schema.from_python_std(t_value, schema_type=cls.schema_type) - return cls(schema) - - @classmethod - def to_flyte_literal_type(cls): - """ - :rtype: flytekit.models.types.LiteralType - """ - return _idl_types.LiteralType(schema=cls.schema_type) - - @classmethod - def promote_from_model(cls, literal_model): - """ - Creates an object of this type from the model primitive defining it. - :param flytekit.models.literals.Literal literal_model: - :rtype: Schema - """ - return cls(_schema_impl.Schema.promote_from_model(literal_model.scalar.schema)) - - @classmethod - def short_class_string(cls): - """ - :rtype: Text - """ - return repr(cls.schema_type) - - def __init__(self, value): - """ - :param flytekit.common.types.impl.schema.Schema value: Schema value to wrap - """ - super(Schema, self).__init__(scalar=_literals.Scalar(schema=value)) - - def to_python_std(self): - """ - :rtype: flytekit.common.types.impl.schema.Schema - """ - return self.scalar.schema - - def short_string(self): - """ - :rtype: Text - """ - return "{}".format( - self.scalar.schema, - ) - - -def schema_instantiator(columns=None): - """ - :param list[(Text, flytekit.common.types.base_sdk_types.FlyteSdkType)] columns: [Optional] Description of the - columns in the underlying schema. Should be tuples with the first element being the name. - :rtype: SchemaInstantiator - """ - if columns is not None and len(columns) == 0: - raise _user_exceptions.FlyteValueException( - columns, - "When specifying a Schema type with a known set of columns, a non-empty list must be provided as " "inputs", - ) - - class _Schema(Schema, metaclass=SchemaInstantiator): - _schema_type = _schema_impl.SchemaType(columns=columns) - - return _Schema - - -def schema_instantiator_from_proto(schema_type): - """ - :param flytekit.models.types.SchemaType schema_type: - :rtype: SchemaInstantiator - """ - - class _Schema(Schema, metaclass=SchemaInstantiator): - _schema_type = _schema_impl.SchemaType.promote_from_model(schema_type) - - return _Schema diff --git a/flytekit/common/workflow.py b/flytekit/common/workflow.py deleted file mode 100644 index 3fc4f498fd..0000000000 --- a/flytekit/common/workflow.py +++ /dev/null @@ -1,309 +0,0 @@ -import datetime as _datetime -from typing import List - -from flytekit.common import constants as _constants -from flytekit.common import interface as _interface -from flytekit.common import nodes as _nodes -from flytekit.common import sdk_bases as _sdk_bases -from flytekit.common.core import identifier as _identifier -from flytekit.common.exceptions import scopes as _exception_scopes -from flytekit.common.exceptions import system as _system_exceptions -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.launch_plan import SdkLaunchPlan -from flytekit.common.mixins import hash as _hash_mixin -from flytekit.common.mixins import registerable as _registerable -from flytekit.configuration import auth as _auth_config -from flytekit.configuration import internal as _internal_config -from flytekit.engines.flyte import engine as _flyte_engine -from flytekit.models import common as _common_models -from flytekit.models import interface as _interface_models -from flytekit.models import launch_plan as _launch_plan_models -from flytekit.models import literals as _literal_models -from flytekit.models import schedule as _schedule_models -from flytekit.models.admin import workflow as _admin_workflow_model -from flytekit.models.core import identifier as _identifier_model -from flytekit.models.core import workflow as _workflow_models - - -class SdkWorkflow( - _hash_mixin.HashOnReferenceMixin, - _registerable.HasDependencies, - _registerable.RegisterableEntity, - _workflow_models.WorkflowTemplate, - metaclass=_sdk_bases.ExtendedSdkType, -): - """ - Previously this class represented both local and control plane constructs. As of this writing, we are making this - class only a control plane class. Workflow constructs that rely on local code being present have been moved to - the new PythonWorkflow class. - """ - - def __init__( - self, - nodes, - interface, - output_bindings, - id, - metadata, - metadata_defaults, - ): - """ - :param list[flytekit.common.nodes.SdkNode] nodes: - :param flytekit.models.interface.TypedInterface interface: Defines a strongly typed interface for the - Workflow (inputs, outputs). This can include some optional parameters. - :param list[flytekit.models.literals.Binding] output_bindings: A list of output bindings that specify how to construct - workflow outputs. Bindings can pull node outputs or specify literals. All workflow outputs specified in - the interface field must be bound - in order for the workflow to be validated. A workflow has an implicit dependency on all of its nodes - to execute successfully in order to bind final outputs. - :param flytekit.models.core.identifier.Identifier id: This is an autogenerated id by the system. The id is - globally unique across Flyte. - :param WorkflowMetadata metadata: This contains information on how to run the workflow. - :param flytekit.models.core.workflow.WorkflowMetadataDefaults metadata_defaults: Defaults to be passed - to nodes contained within workflow. - """ - for n in nodes: - for upstream in n.upstream_nodes: - if upstream.id is None: - raise _user_exceptions.FlyteAssertion( - "Some nodes contained in the workflow were not found in the workflow description. Please " - "ensure all nodes are either assigned to attributes within the class or an element in a " - "list, dict, or tuple which is stored as an attribute in the class." - ) - - super(SdkWorkflow, self).__init__( - id=id, - metadata=metadata, - metadata_defaults=metadata_defaults, - interface=interface, - nodes=nodes, - outputs=output_bindings, - ) - self._sdk_nodes = nodes - self._has_registered = False - - @property - def upstream_entities(self): - return set(n.executable_sdk_object for n in self._sdk_nodes) - - @property - def interface(self): - """ - :rtype: flytekit.common.interface.TypedInterface - """ - return super(SdkWorkflow, self).interface - - @property - def entity_type_text(self): - """ - :rtype: Text - """ - return "Workflow" - - @property - def resource_type(self): - """ - Integer from _identifier.ResourceType enum - :rtype: int - """ - return _identifier_model.ResourceType.WORKFLOW - - def get_sub_workflows(self): - """ - Recursive call that returns all subworkflows in the current workflow - - :rtype: list[SdkWorkflow] - """ - result = [] - for node in self.nodes: - if node.workflow_node is not None and node.workflow_node.sub_workflow_ref is not None: - if node.executable_sdk_object is not None and node.executable_sdk_object.entity_type_text == "Workflow": - result.append(node.executable_sdk_object) - result.extend(node.executable_sdk_object.get_sub_workflows()) - else: - raise _system_exceptions.FlyteSystemException( - "workflow node with subworkflow found but bad executable " - "object {}".format(node.executable_sdk_object) - ) - - # get subworkflows in conditional branches - if node.branch_node is not None: - if_else: _workflow_models.IfElseBlock = node.branch_node.if_else - leaf_nodes: List[_nodes.SdkNode] = filter( - None, - [ - if_else.case.then_node, - *([] if if_else.other is None else [x.then_node for x in if_else.other]), - if_else.else_node, - ], - ) - for leaf_node in leaf_nodes: - exec_sdk_obj = leaf_node.executable_sdk_object - if exec_sdk_obj is not None and exec_sdk_obj.entity_type_text == "Workflow": - result.append(exec_sdk_obj) - result.extend(exec_sdk_obj.get_sub_workflows()) - - return result - - @classmethod - @_exception_scopes.system_entry_point - def fetch(cls, project, domain, name, version=None): - """ - This function uses the engine loader to call create a hydrated task from Admin. - :param Text project: - :param Text domain: - :param Text name: - :param Text version: - :rtype: SdkWorkflow - """ - version = version or _internal_config.VERSION.get() - workflow_id = _identifier.Identifier(_identifier_model.ResourceType.WORKFLOW, project, domain, name, version) - admin_workflow = _flyte_engine.get_client().get_workflow(workflow_id) - cwc = admin_workflow.closure.compiled_workflow - primary_template = cwc.primary.template - sub_workflow_map = {sw.template.id: sw.template for sw in cwc.sub_workflows} - task_map = {t.template.id: t.template for t in cwc.tasks} - sdk_workflow = cls.promote_from_model(primary_template, sub_workflow_map, task_map) - sdk_workflow._id = workflow_id - sdk_workflow._has_registered = True - return sdk_workflow - - @classmethod - def get_non_system_nodes(cls, nodes): - """ - :param list[flytekit.models.core.workflow.Node] nodes: - :rtype: list[flytekit.models.core.workflow.Node] - """ - return [n for n in nodes if n.id not in {_constants.START_NODE_ID, _constants.END_NODE_ID}] - - @classmethod - def promote_from_model(cls, base_model, sub_workflows=None, tasks=None): - """ - :param flytekit.models.core.workflow.WorkflowTemplate base_model: - :param dict[flytekit.models.core.identifier.Identifier, flytekit.models.core.workflow.WorkflowTemplate] - sub_workflows: Provide a list of WorkflowTemplate - models (should be returned from Admin as part of the admin CompiledWorkflowClosure. Relevant sub-workflows - should always be provided. - :param dict[flytekit.models.core.identifier.Identifier, flytekit.models.task.TaskTemplate] tasks: Same as above - but for tasks. If tasks are not provided relevant TaskTemplates will be fetched from Admin - :rtype: SdkWorkflow - """ - base_model_non_system_nodes = cls.get_non_system_nodes(base_model.nodes) - sub_workflows = sub_workflows or {} - tasks = tasks or {} - node_map = { - n.id: _nodes.SdkNode.promote_from_model(n, sub_workflows, tasks) for n in base_model_non_system_nodes - } - - # Set upstream nodes for each node - for n in base_model_non_system_nodes: - current = node_map[n.id] - for upstream_id in current.upstream_node_ids: - upstream_node = node_map[upstream_id] - current << upstream_node - - # No inputs/outputs specified, see the constructor for more information on the overrides. - return cls( - nodes=list(node_map.values()), - id=_identifier.Identifier.promote_from_model(base_model.id), - metadata=base_model.metadata, - metadata_defaults=base_model.metadata_defaults, - interface=_interface.TypedInterface.promote_from_model(base_model.interface), - output_bindings=base_model.outputs, - ) - - @_exception_scopes.system_entry_point - def register(self, project, domain, name, version): - """ - :param Text project: - :param Text domain: - :param Text name: - :param Text version: - """ - self.validate() - id_to_register = _identifier.Identifier(_identifier_model.ResourceType.WORKFLOW, project, domain, name, version) - old_id = self.id - self._id = id_to_register - try: - client = _flyte_engine.get_client() - sub_workflows = self.get_sub_workflows() - client.create_workflow( - id_to_register, - _admin_workflow_model.WorkflowSpec( - self, - sub_workflows, - ), - ) - self._id = id_to_register - self._has_registered = True - return str(id_to_register) - except _user_exceptions.FlyteEntityAlreadyExistsException: - pass - except Exception: - self._id = old_id - raise - - @_exception_scopes.system_entry_point - def serialize(self): - """ - Serializing a workflow should produce an object similar to what the registration step produces, in preparation - for actual registration to Admin. - - :rtype: flyteidl.admin.workflow_pb2.WorkflowSpec - """ - sub_workflows = self.get_sub_workflows() - return _admin_workflow_model.WorkflowSpec( - self, - sub_workflows, - ).to_flyte_idl() - - @_exception_scopes.system_entry_point - def validate(self): - pass - - @_exception_scopes.system_entry_point - def create_launch_plan(self, *args, **kwargs): - # TODO: Correct after implementing new launch plan - assumable_iam_role = _auth_config.ASSUMABLE_IAM_ROLE.get() - kubernetes_service_account = _auth_config.KUBERNETES_SERVICE_ACCOUNT.get() - - if not (assumable_iam_role or kubernetes_service_account): - raise _user_exceptions.FlyteValidationException("No assumable role or service account found") - auth_role = _common_models.AuthRole( - assumable_iam_role=assumable_iam_role, - kubernetes_service_account=kubernetes_service_account, - ) - - return SdkLaunchPlan( - workflow_id=self.id, - entity_metadata=_launch_plan_models.LaunchPlanMetadata( - schedule=_schedule_models.Schedule(""), - notifications=[], - ), - default_inputs=_interface_models.ParameterMap({}), - fixed_inputs=_literal_models.LiteralMap(literals={}), - labels=_common_models.Labels({}), - annotations=_common_models.Annotations({}), - auth_role=auth_role, - raw_output_data_config=_common_models.RawOutputDataConfig(""), - ) - - @_exception_scopes.system_entry_point - def __call__(self, *args, **input_map): - if len(args) > 0: - raise _user_exceptions.FlyteAssertion( - "When adding a workflow as a node in a workflow, all inputs must be specified with kwargs only. We " - "detected {} positional args.".format(len(args)) - ) - bindings, upstream_nodes = self.interface.create_bindings_for_inputs(input_map) - - node = _nodes.SdkNode( - id=None, - metadata=_workflow_models.NodeMetadata( - "placeholder", _datetime.timedelta(), _literal_models.RetryStrategy(0) - ), - upstream_nodes=upstream_nodes, - bindings=sorted(bindings, key=lambda b: b.var), - sdk_workflow=self, - ) - return node diff --git a/flytekit/common/workflow_execution.py b/flytekit/common/workflow_execution.py deleted file mode 100644 index 14695d0e68..0000000000 --- a/flytekit/common/workflow_execution.py +++ /dev/null @@ -1,183 +0,0 @@ -import os as _os - -import six as _six -from flyteidl.core import literals_pb2 as _literals_pb2 - -from flytekit.clients.helpers import iterate_node_executions as _iterate_node_executions -from flytekit.common import nodes as _nodes -from flytekit.common import sdk_bases as _sdk_bases -from flytekit.common import utils as _common_utils -from flytekit.common.core import identifier as _core_identifier -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.mixins import artifact as _artifact -from flytekit.common.types import helpers as _type_helpers -from flytekit.engines.flyte import engine as _flyte_engine -from flytekit.interfaces.data import data_proxy as _data_proxy -from flytekit.models import execution as _execution_models -from flytekit.models import literals as _literal_models -from flytekit.models.core import execution as _core_execution_models - - -class SdkWorkflowExecution( - _execution_models.Execution, - _artifact.ExecutionArtifact, - metaclass=_sdk_bases.ExtendedSdkType, -): - def __init__(self, *args, **kwargs): - super(SdkWorkflowExecution, self).__init__(*args, **kwargs) - self._node_executions = None - self._inputs = None - self._outputs = None - - @property - def node_executions(self): - """ - :rtype: dict[Text, flytekit.common.nodes.SdkNodeExecution] - """ - return self._node_executions or {} - - @property - def inputs(self): - """ - Returns the inputs to the execution in the standard Python format as dictated by the type engine. - :rtype: dict[Text, T] - """ - if self._inputs is None: - client = _flyte_engine.get_client() - execution_data = client.get_execution_data(self.id) - - # Inputs are returned inline unless they are too big, in which case a url blob pointing to them is returned. - if bool(execution_data.full_inputs.literals): - input_map = execution_data.full_inputs - elif execution_data.inputs.bytes > 0: - with _common_utils.AutoDeletingTempDir() as t: - tmp_name = _os.path.join(t.name, "inputs.pb") - _data_proxy.Data.get_data(execution_data.inputs.url, tmp_name) - input_map = _literal_models.LiteralMap.from_flyte_idl( - _common_utils.load_proto_from_file(_literals_pb2.LiteralMap, tmp_name) - ) - else: - input_map = _literal_models.LiteralMap({}) - - self._inputs = _type_helpers.unpack_literal_map_to_sdk_python_std(input_map) - return self._inputs - - @property - def outputs(self): - """ - Returns the outputs to the execution in the standard Python format as dictated by the type engine. If the - execution ended in error or the execution is in progress, an exception will be raised. - :rtype: dict[Text, T] or None - """ - if not self.is_complete: - raise _user_exceptions.FlyteAssertion( - "Please what until the node execution has completed before " "requesting the outputs." - ) - if self.error: - raise _user_exceptions.FlyteAssertion("Outputs could not be found because the execution ended in failure.") - - if self._outputs is None: - client = _flyte_engine.get_client() - - execution_data = client.get_execution_data(self.id) - # Outputs are returned inline unless they are too big, in which case a url blob pointing to them is returned. - if bool(execution_data.full_outputs.literals): - output_map = execution_data.full_outputs - - elif execution_data.outputs.bytes > 0: - with _common_utils.AutoDeletingTempDir() as t: - tmp_name = _os.path.join(t.name, "outputs.pb") - _data_proxy.Data.get_data(execution_data.outputs.url, tmp_name) - output_map = _literal_models.LiteralMap.from_flyte_idl( - _common_utils.load_proto_from_file(_literals_pb2.LiteralMap, tmp_name) - ) - else: - output_map = _literal_models.LiteralMap({}) - - self._outputs = _type_helpers.unpack_literal_map_to_sdk_python_std(output_map) - return self._outputs - - @property - def error(self): - """ - If execution is in progress, raise an exception. Otherwise, return None if no error was present upon - reaching completion. - :rtype: flytekit.models.core.execution.ExecutionError or None - """ - if not self.is_complete: - raise _user_exceptions.FlyteAssertion( - "Please wait until a workflow has completed before checking for an " "error." - ) - return self.closure.error - - @property - def is_complete(self): - """ - Dictates whether or not the execution is complete. - :rtype: bool - """ - return self.closure.phase in { - _core_execution_models.WorkflowExecutionPhase.ABORTED, - _core_execution_models.WorkflowExecutionPhase.FAILED, - _core_execution_models.WorkflowExecutionPhase.SUCCEEDED, - _core_execution_models.WorkflowExecutionPhase.TIMED_OUT, - } - - @classmethod - def promote_from_model(cls, base_model): - """ - :param _execution_models.Execution base_model: - :rtype: SdkWorkflowExecution - """ - return cls( - closure=base_model.closure, - id=_core_identifier.WorkflowExecutionIdentifier.promote_from_model(base_model.id), - spec=base_model.spec, - ) - - @classmethod - def fetch(cls, project, domain, name): - """ - :param Text project: - :param Text domain: - :param Text name: - :rtype: SdkWorkflowExecution - """ - wf_exec_id = _core_identifier.WorkflowExecutionIdentifier(project=project, domain=domain, name=name) - admin_exec = _flyte_engine.get_client().get_execution(wf_exec_id) - - return cls.promote_from_model(admin_exec) - - def sync(self): - """ - Syncs the state of the underlying execution artifact with the state observed by the platform. - :rtype: None - """ - if not self.is_complete or self._node_executions is None: - self._sync_closure() - self._node_executions = self.get_node_executions() - - def _sync_closure(self): - """ - Syncs the closure of the underlying execution artifact with the state observed by the platform. - :rtype: None - """ - if not self.is_complete: - client = _flyte_engine.get_client() - self._closure = client.get_execution(self.id).closure - - def get_node_executions(self, filters=None): - """ - :param list[flytekit.models.filters.Filter] filters: - :rtype: dict[Text, flytekit.common.nodes.SdkNodeExecution] - """ - client = _flyte_engine.get_client() - node_exec_models = {v.id.node_id: v for v in _iterate_node_executions(client, self.id, filters=filters)} - - return {k: _nodes.SdkNodeExecution.promote_from_model(v) for k, v in _six.iteritems(node_exec_models)} - - def terminate(self, cause): - """ - :param Text cause: - """ - _flyte_engine.get_client().terminate_execution(self.id, cause) diff --git a/flytekit/configuration/__init__.py b/flytekit/configuration/__init__.py index f6fc2c633e..6dd9c28d35 100644 --- a/flytekit/configuration/__init__.py +++ b/flytekit/configuration/__init__.py @@ -1,12 +1,6 @@ import logging as _logging import os as _os - -import six as _six - -try: - import pathlib as _pathlib -except ImportError: - import pathlib2 as _pathlib # python 2 backport +import pathlib as _pathlib def set_flyte_config_file(config_file_path): @@ -38,7 +32,7 @@ def __init__(self, new_config_path, internal_overrides=None): import flytekit.configuration.common as _common self._internal_overrides = { - _common.format_section_key("internal", k): v for k, v in _six.iteritems(internal_overrides or {}) + _common.format_section_key("internal", k): v for k, v in (internal_overrides or {}).items() } self._new_config_path = new_config_path self._old_config_path = None @@ -47,13 +41,13 @@ def __init__(self, new_config_path, internal_overrides=None): def __enter__(self): import flytekit.configuration.internal as _internal - self._old_internals = {k: _os.environ.get(k) for k in _six.iterkeys(self._internal_overrides)} + self._old_internals = {k: _os.environ.get(k) for k in self._internal_overrides.keys()} self._old_config_path = _os.environ.get(_internal.CONFIGURATION_PATH.env_var) _os.environ.update(self._internal_overrides) set_flyte_config_file(self._new_config_path) def __exit__(self, exc_type, exc_val, exc_tb): - for k, v in _six.iteritems(self._old_internals): + for k, v in self._old_internals.items(): if v is not None: _os.environ[k] = v else: diff --git a/flytekit/configuration/common.py b/flytekit/configuration/common.py index 93f5c2d403..6e0b2088a4 100644 --- a/flytekit/configuration/common.py +++ b/flytekit/configuration/common.py @@ -2,7 +2,7 @@ import configparser as _configparser import os as _os -from flytekit.common.exceptions import user as _user_exceptions +from flytekit.exceptions import user as _user_exceptions def format_section_key(section, key): diff --git a/flytekit/configuration/platform.py b/flytekit/configuration/platform.py index d420e7a019..5c4061fa4f 100644 --- a/flytekit/configuration/platform.py +++ b/flytekit/configuration/platform.py @@ -1,5 +1,5 @@ -from flytekit.common import constants as _constants from flytekit.configuration import common as _config_common +from flytekit.core import constants as _constants URL = _config_common.FlyteStringConfigurationEntry("platform", "url") diff --git a/flytekit/contrib/__init__.py b/flytekit/contrib/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/flytekit/contrib/sensors/__init__.py b/flytekit/contrib/sensors/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/flytekit/contrib/sensors/base_sensor.py b/flytekit/contrib/sensors/base_sensor.py deleted file mode 100644 index 8522918a34..0000000000 --- a/flytekit/contrib/sensors/base_sensor.py +++ /dev/null @@ -1,90 +0,0 @@ -import abc as _abc -import datetime as _datetime -import logging as _logging -import sys as _sys -import time as _time -import traceback as _traceback - -import six as _six - - -class Sensor(object, metaclass=_abc.ABCMeta): - def __init__(self, evaluation_interval=None, max_failures=0): - """ - :param datetime.timedelta evaluation_interval: This is the time to wait between evaluation attempts of this - sensor. If the sensor takes longer to evaluate than the poll_interval, it will immediately begin - evaluation again. - :param int max_failures: This is the maximum number of failures that can happen while attempting to sense - before perma-failing. - """ - if evaluation_interval is None: - evaluation_interval = _datetime.timedelta(seconds=30) - self._evaluation_interval = evaluation_interval - self._max_failures = max_failures - self._failures = 0 - self._exc_info = None - self._last_executed_time = _datetime.datetime(year=1990, month=6, day=30) # Arbitrary date in the past. - self._sensed = False - - @_abc.abstractmethod - def _do_poll(self): - """ - :rtype: (bool, Optional[datetime.timedelta]) - """ - pass - - def sense_with_wait_hint(self): - """ - Attempts to sense based on the lambda expression. The method will return the last sensed result. If the - rate of sensing is exceeded for the sensor, the timedelta in the returned tuple will tell the caller how long it - should sleep before trying again. - :rtype: (bool, Optional[datetime.timedelta]) - """ - # Return cached success. This simplifies code for the conditional sensors. - if self._sensed: - return self._sensed, self._evaluation_interval - - # Perma-fail to prevent abuse of sensed objects. - if self._failures > self._max_failures: - _six.reraise(*self._exc_info) - - now = _datetime.datetime.utcnow() - - time_to_wait_eval_period = self._evaluation_interval - (now - self._last_executed_time) - if time_to_wait_eval_period > _datetime.timedelta(): - return self._sensed, time_to_wait_eval_period - - try: - self._sensed, time_to_wait = self._do_poll() - time_to_wait = time_to_wait or self._evaluation_interval - except BaseException: - self._failures += 1 - self._exc_info = _sys.exc_info() - if self._failures > self._max_failures: - _logging.error( - "{} failed (with no remaining retries) due to:\n\n{}".format(self, _traceback.format_exc()), - ) - raise - else: - _logging.warn("{} failed (but will retry) due to:\n\n{}".format(self, _traceback.format_exc())) - time_to_wait = self._evaluation_interval - - self._last_executed_time = _datetime.datetime.utcnow() - - return self._sensed, time_to_wait - - def sense(self, timeout=None): - """ - Attempts - :param datetime.timedelta timeout: - :rtype: bool - """ - started = _datetime.datetime.utcnow() - while True: - sensed, time_to_wait = self.sense_with_wait_hint() - if sensed: - return True - if time_to_wait: - _time.sleep(time_to_wait.total_seconds()) - if timeout is not None and (_datetime.datetime.utcnow() - started) > timeout: - return False diff --git a/flytekit/contrib/sensors/impl.py b/flytekit/contrib/sensors/impl.py deleted file mode 100644 index 8276320922..0000000000 --- a/flytekit/contrib/sensors/impl.py +++ /dev/null @@ -1,107 +0,0 @@ -from flytekit.contrib.sensors.base_sensor import Sensor as _Sensor -from flytekit.plugins import hmsclient as _hmsclient - - -class _HiveSensor(_Sensor): - def __init__(self, host, port, schema="default", **kwargs): - """ - :param Text host: - :param Text port: - :param Text schema: The schema/database that we should consider. - :param **kwargs: See flytekit.contrib.sensors.base_sensor.Sensor for more - parameters. - """ - self._schema = schema - self._host = host - self._port = port - self._hive_metastore_client = _hmsclient.HMSClient(host=host, port=port) - super(_HiveSensor, self).__init__(**kwargs) - - -class HiveTableSensor(_HiveSensor): - def __init__(self, table_name, host, port, **kwargs): - """ - :param Text host: The host for the Hive metastore Thrift service. - :param Text port: The port for the Hive metastore Thrift Service - :param Text table_name: The name of the table to look for. - :param **kwargs: See _HiveSensor and flytekit.contrib.sensors.base_sensor.Sensor for more - parameters. - """ - super(HiveTableSensor, self).__init__(host, port, **kwargs) - self._table_name = table_name - - def _do_poll(self): - """ - :rtype: (bool, Optional[datetime.timedelta]) - """ - with self._hive_metastore_client as client: - try: - client.get_table(self._schema, self._table_name) - return True, None - except _hmsclient.genthrift.hive_metastore.ttypes.NoSuchObjectException: - return False, None - - -class HiveNamedPartitionSensor(_HiveSensor): - def __init__(self, table_name, partition_names, host, port, **kwargs): - """ - This class allows sensing for a specific named Hive Partition. This is the preferred partition sensing - operator because it is more efficient than evaluating a filter expression. - - :param Text table_name: The name of the table - :param Text partition_name: The name of the partition to listen for (example: 'ds=2017-01-01/region=NYC') - :param Text host: The host for the Hive metastore Thrift service. - :param Text port: The port for the Hive metastore Thrift Service - :param **kwargs: See _HiveSensor and flytekit.contrib.sensors.base_sensor.Sensor for more - parameters. - """ - super(HiveNamedPartitionSensor, self).__init__(host, port, **kwargs) - self._table_name = table_name - self._partition_names = partition_names - - def _do_poll(self): - """ - :rtype: (bool, Optional[datetime.timedelta]) - """ - with self._hive_metastore_client as client: - try: - for partition_name in self._partition_names: - client.get_partition_by_name(self._schema, self._table_name, partition_name) - return True, None - except _hmsclient.genthrift.hive_metastore.ttypes.NoSuchObjectException: - return False, None - - -class HiveFilteredPartitionSensor(_HiveSensor): - def __init__(self, table_name, partition_filter, host, port, **kwargs): - """ - This class allows sensing for any Hive partition that matches a filter expression. It is recommended that the - user should use HiveNamedPartitionSensor instead when possible because it is a more efficient API. - - :param Text table_name: The name of the table - :param Text partition_filter: A filter expression for the partition. (example: "ds = '2017-01-01' and - region='NYC') - :param Text host: The host for the Hive metastore Thrift service. - :param Text port: The port for the Hive metastore Thrift Service - :param **kwargs: See _HiveSensor and flytekit.contrib.sensors.base_sensor.Sensor for more - parameters. - """ - super(HiveFilteredPartitionSensor, self).__init__(host, port, **kwargs) - self._table_name = table_name - self._partition_filter = partition_filter - - def _do_poll(self): - """ - :rtype: (bool, Optional[datetime.timedelta]) - """ - with self._hive_metastore_client as client: - partitions = client.get_partitions_by_filter( - db_name=self._schema, - tbl_name=self._table_name, - filter=self._partition_filter, - max_parts=1, - ) - if partitions: - return True, None - else: - return False, None diff --git a/flytekit/contrib/sensors/task.py b/flytekit/contrib/sensors/task.py deleted file mode 100644 index 0749fc39dc..0000000000 --- a/flytekit/contrib/sensors/task.py +++ /dev/null @@ -1,127 +0,0 @@ -from flytekit.common import constants as _common_constants -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.tasks import sdk_runnable as _sdk_runnable -from flytekit.contrib.sensors.base_sensor import Sensor as _Sensor - - -class SensorTask(_sdk_runnable.SdkRunnableTask): - def _execute_user_code(self, context, inputs): - sensor = super(SensorTask, self)._execute_user_code(context=context, inputs=inputs) - if sensor is not None: - if not isinstance(sensor, _Sensor): - raise _user_exceptions.FlyteTypeException( - received_type=type(sensor), - expected_type=_Sensor, - ) - succeeded = sensor.sense() - if not succeeded: - raise _user_exceptions.FlyteRecoverableException() - - -def sensor_task( - _task_function=None, - retries=0, - interruptible=None, - deprecated="", - storage_request=None, - cpu_request=None, - gpu_request=None, - memory_request=None, - storage_limit=None, - cpu_limit=None, - gpu_limit=None, - memory_limit=None, - timeout=None, - environment=None, - cls=None, -): - """ - Decorator to create a Sensor Task definition. This task will run as a single unit of work on the platform. - - .. code-block:: python - @sensor_task(retries=3) - def my_task(wf_params): - return HiveTableSensor( - schema='default', - table_name='mocked_table', - host='localhost', - port=1234, - ) - - :param _task_function: this is the decorated method and shouldn't be declared explicitly. The function must - take a first argument, and then named arguments matching those defined in @inputs. No keyword - arguments are allowed for wrapped task functions. - :param int retries: [optional] integer determining number of times task can be retried on - :py:exc:`flytekit.sdk.exceptions.RecoverableException` or transient platform failures. Defaults - to 0. - .. note:: - If retries > 0, the task must be able to recover from any remote state created within the user code. It is - strongly recommended that tasks are written to be idempotent. - :param bool interruptible: Specify whether task is interruptible - :param Text deprecated: [optional] string that should be provided if this task is deprecated. The string - will be logged as a warning so it should contain information regarding how to update to a newer task. - :param Text storage_request: [optional] Kubernetes resource string for lower-bound of disk storage space - for the task to run. Default is set by platform-level configuration. - .. note:: - This is currently not supported by the platform. - :param Text cpu_request: [optional] Kubernetes resource string for lower-bound of cores for the task to execute. - This can be set to a fractional portion of a CPU. Default is set by platform-level configuration. - TODO: Add links to resource string documentation for Kubernetes - :param Text gpu_request: [optional] Kubernetes resource string for lower-bound of desired GPUs. - Default is set by platform-level configuration. - TODO: Add links to resource string documentation for Kubernetes - :param Text memory_request: [optional] Kubernetes resource string for lower-bound of physical memory - necessary for the task to execute. Default is set by platform-level configuration. - TODO: Add links to resource string documentation for Kubernetes - :param Text storage_limit: [optional] Kubernetes resource string for upper-bound of disk storage space - for the task to run. This amount is not guaranteed! If not specified, it is set equal to storage_request. - .. note:: - This is currently not supported by the platform. - :param Text cpu_limit: [optional] Kubernetes resource string for upper-bound of cores for the task to execute. - This can be set to a fractional portion of a CPU. This amount is not guaranteed! If not specified, - it is set equal to cpu_request. - :param Text gpu_limit: [optional] Kubernetes resource string for upper-bound of desired GPUs. This amount is not - guaranteed! If not specified, it is set equal to gpu_request. - :param Text memory_limit: [optional] Kubernetes resource string for upper-bound of physical memory - necessary for the task to execute. This amount is not guaranteed! If not specified, it is set equal to - memory_request. - :param datetime.timedelta timeout: [optional] describes how long the task should be allowed to - run at max before triggering a retry (if retries are enabled). By default, tasks are allowed to run - indefinitely. If a null timedelta is passed (i.e. timedelta(seconds=0)), the task will not timeout. - :param dict[Text,Text] environment: [optional] environment variables to set when executing this task. - :param cls: This can be used to override the task implementation with a user-defined extension. The class - provided must be a subclass of flytekit.common.tasks.sdk_runnable.SdkRunnableTask. A user can use this to - inject bespoke logic into the base Flyte programming model. Ideally, should be a sub-class of SensorTask or - otherwise mimic the behavior. - :rtype: SensorTask - """ - - def wrapper(fn): - return (SensorTask or cls)( - task_function=fn, - task_type=_common_constants.SdkTaskType.SENSOR_TASK, - retries=retries, - interruptible=interruptible, - deprecated=deprecated, - storage_request=storage_request, - cpu_request=cpu_request, - gpu_request=gpu_request, - memory_request=memory_request, - storage_limit=storage_limit, - cpu_limit=cpu_limit, - gpu_limit=gpu_limit, - memory_limit=memory_limit, - timeout=timeout, - environment=environment, - custom={}, - discovery_version="", - discoverable=False, - cache_serializable=False, - ) - - # This is syntactic-sugar, so that when calling this decorator without args, you can either - # do it with () or without any () - if _task_function: - return wrapper(_task_function) - else: - return wrapper diff --git a/flytekit/core/base_task.py b/flytekit/core/base_task.py index 53fa011185..f590ba2033 100644 --- a/flytekit/core/base_task.py +++ b/flytekit/core/base_task.py @@ -23,8 +23,13 @@ from dataclasses import dataclass from typing import Any, Dict, Generic, List, Optional, Tuple, Type, TypeVar, Union -from flytekit.common.tasks.sdk_runnable import ExecutionParameters -from flytekit.core.context_manager import FlyteContext, FlyteContextManager, FlyteEntities, SerializationSettings +from flytekit.core.context_manager import ( + ExecutionParameters, + FlyteContext, + FlyteContextManager, + FlyteEntities, + SerializationSettings, +) from flytekit.core.interface import Interface, transform_interface_to_typed_interface from flytekit.core.local_cache import LocalTaskCache from flytekit.core.promise import ( diff --git a/flytekit/common/constants.py b/flytekit/core/constants.py similarity index 100% rename from flytekit/common/constants.py rename to flytekit/core/constants.py diff --git a/flytekit/core/container_task.py b/flytekit/core/container_task.py index 75c6449c07..d46057c623 100644 --- a/flytekit/core/container_task.py +++ b/flytekit/core/container_task.py @@ -1,11 +1,11 @@ from enum import Enum from typing import Any, Dict, List, Optional, Type -from flytekit.common.tasks.raw_container import _get_container_definition from flytekit.core.base_task import PythonTask, TaskMetadata from flytekit.core.context_manager import SerializationSettings from flytekit.core.interface import Interface from flytekit.core.resources import Resources, ResourceSpec +from flytekit.core.utils import _get_container_definition from flytekit.models import task as _task_model diff --git a/flytekit/core/context_manager.py b/flytekit/core/context_manager.py index 627cf3ea11..1db5e11d5e 100644 --- a/flytekit/core/context_manager.py +++ b/flytekit/core/context_manager.py @@ -23,19 +23,21 @@ import typing from contextlib import contextmanager from dataclasses import dataclass, field +from datetime import datetime from enum import Enum from typing import Any, Dict, Generator, List, Optional, Union from docker_image import reference from flytekit.clients import friendly as friendly_client # noqa -from flytekit.common.core.identifier import WorkflowExecutionIdentifier as _SdkWorkflowExecutionIdentifier -from flytekit.common.tasks.sdk_runnable import ExecutionParameters from flytekit.configuration import images, internal from flytekit.configuration import sdk as _sdk_config +from flytekit.configuration import secrets +from flytekit.core import mock_stats, utils from flytekit.core.data_persistence import FileAccessProvider, default_local_file_access_provider from flytekit.core.node import Node -from flytekit.engines.unit import mock_stats as _mock_stats +from flytekit.interfaces.cli_identifiers import WorkflowExecutionIdentifier +from flytekit.interfaces.stats import taggable from flytekit.models.core import identifier as _identifier # TODO: resolve circular import from flytekit.core.python_auto_container import TaskResolverMixin @@ -132,6 +134,216 @@ def get_image_config(img_name: Optional[str] = None) -> ImageConfig: return ImageConfig(default_image=default_img, images=other_images) +class ExecutionParameters(object): + """ + This is a run-time user-centric context object that is accessible to every @task method. It can be accessed using + + .. code-block:: python + + flytekit.current_context() + + This object provides the following + * a statsd handler + * a logging handler + * the execution ID as an :py:class:`flytekit.models.core.identifier.WorkflowExecutionIdentifier` object + * a working directory for the user to write arbitrary files to + + Please do not confuse this object with the :py:class:`flytekit.FlyteContext` object. + """ + + @dataclass(init=False) + class Builder(object): + stats: taggable.TaggableStats + execution_date: datetime + logging: _logging + execution_id: str + attrs: typing.Dict[str, typing.Any] + working_dir: typing.Union[os.PathLike, utils.AutoDeletingTempDir] + + def __init__(self, current: typing.Optional[ExecutionParameters] = None): + self.stats = current.stats if current else None + self.execution_date = current.execution_date if current else None + self.working_dir = current.working_directory if current else None + self.execution_id = current.execution_id if current else None + self.logging = current.logging if current else None + self.attrs = current._attrs if current else {} + + def add_attr(self, key: str, v: typing.Any) -> ExecutionParameters.Builder: + self.attrs[key] = v + return self + + def build(self) -> ExecutionParameters: + if not isinstance(self.working_dir, utils.AutoDeletingTempDir): + pathlib.Path(self.working_dir).mkdir(parents=True, exist_ok=True) + return ExecutionParameters( + execution_date=self.execution_date, + stats=self.stats, + tmp_dir=self.working_dir, + execution_id=self.execution_id, + logging=self.logging, + **self.attrs, + ) + + @staticmethod + def new_builder(current: ExecutionParameters = None) -> Builder: + return ExecutionParameters.Builder(current=current) + + def builder(self) -> Builder: + return ExecutionParameters.Builder(current=self) + + def __init__(self, execution_date, tmp_dir, stats, execution_id, logging, **kwargs): + """ + Args: + execution_date: Date when the execution is running + tmp_dir: temporary directory for the execution + stats: handle to emit stats + execution_id: Identifier for the xecution + logging: handle to logging + """ + self._stats = stats + self._execution_date = execution_date + self._working_directory = tmp_dir + self._execution_id = execution_id + self._logging = logging + # AutoDeletingTempDir's should be used with a with block, which creates upon entry + self._attrs = kwargs + # It is safe to recreate the Secrets Manager + self._secrets_manager = SecretsManager() + + @property + def stats(self) -> taggable.TaggableStats: + """ + A handle to a special statsd object that provides usefully tagged stats. + TODO: Usage examples and better comments + """ + return self._stats + + @property + def logging(self) -> _logging: + """ + A handle to a useful logging object. + TODO: Usage examples + """ + return self._logging + + @property + def working_directory(self) -> utils.AutoDeletingTempDir: + """ + A handle to a special working directory for easily producing temporary files. + + TODO: Usage examples + TODO: This does not always return a AutoDeletingTempDir + """ + return self._working_directory + + @property + def execution_date(self) -> datetime: + """ + This is a datetime representing the time at which a workflow was started. This is consistent across all tasks + executed in a workflow or sub-workflow. + + .. note:: + + Do NOT use this execution_date to drive any production logic. It might be useful as a tag for data to help + in debugging. + """ + return self._execution_date + + @property + def execution_id(self) -> str: + """ + This is the identifier of the workflow execution within the underlying engine. It will be consistent across all + task executions in a workflow or sub-workflow execution. + + .. note:: + + Do NOT use this execution_id to drive any production logic. This execution ID should only be used as a tag + on output data to link back to the workflow run that created it. + """ + return self._execution_id + + @property + def secrets(self) -> SecretsManager: + return self._secrets_manager + + def __getattr__(self, attr_name: str) -> typing.Any: + """ + This houses certain task specific context. For example in Spark, it houses the SparkSession, etc + """ + attr_name = attr_name.upper() + if self._attrs and attr_name in self._attrs: + return self._attrs[attr_name] + raise AssertionError(f"{attr_name} not available as a parameter in Flyte context - are you in right task-type?") + + def has_attr(self, attr_name: str) -> bool: + attr_name = attr_name.upper() + if self._attrs and attr_name in self._attrs: + return True + return False + + def get(self, key: str) -> typing.Any: + """ + Returns task specific context if present else raise an error. The returned context will match the key + """ + return self.__getattr__(attr_name=key) + + +class SecretsManager(object): + """ + This provides a secrets resolution logic at runtime. + The resolution order is + - Try env var first. The env var should have the configuration.SECRETS_ENV_PREFIX. The env var will be all upper + cased + - If not then try the file where the name matches lower case + ``configuration.SECRETS_DEFAULT_DIR//configuration.SECRETS_FILE_PREFIX`` + + All configuration values can always be overridden by injecting an environment variable + """ + + def __init__(self): + self._base_dir = str(secrets.SECRETS_DEFAULT_DIR.get()).strip() + self._file_prefix = str(secrets.SECRETS_FILE_PREFIX.get()).strip() + self._env_prefix = str(secrets.SECRETS_ENV_PREFIX.get()).strip() + + def get(self, group: str, key: str) -> str: + """ + Retrieves a secret using the resolution order -> Env followed by file. If not found raises a ValueError + """ + self.check_group_key(group, key) + env_var = self.get_secrets_env_var(group, key) + fpath = self.get_secrets_file(group, key) + v = os.environ.get(env_var) + if v is not None: + return v + if os.path.exists(fpath): + with open(fpath, "r") as f: + return f.read().strip() + raise ValueError( + f"Unable to find secret for key {key} in group {group} " f"in Env Var:{env_var} and FilePath: {fpath}" + ) + + def get_secrets_env_var(self, group: str, key: str) -> str: + """ + Returns a string that matches the ENV Variable to look for the secrets + """ + self.check_group_key(group, key) + return f"{self._env_prefix}{group.upper()}_{key.upper()}" + + def get_secrets_file(self, group: str, key: str) -> str: + """ + Returns a path that matches the file to look for the secrets + """ + self.check_group_key(group, key) + return os.path.join(self._base_dir, group.lower(), f"{self._file_prefix}{key.lower()}") + + @staticmethod + def check_group_key(group: str, key: str): + if group is None or group == "": + raise ValueError("secrets group is a mandatory field.") + if key is None or key == "": + raise ValueError("secrets key is a mandatory field.") + + @dataclass class EntrypointSettings(object): """ @@ -684,9 +896,9 @@ def initialize(): # Note we use the SdkWorkflowExecution object purely for formatting into the ex:project:domain:name format users # are already acquainted with default_user_space_params = ExecutionParameters( - execution_id=str(_SdkWorkflowExecutionIdentifier.promote_from_model(default_execution_id)), + execution_id=str(WorkflowExecutionIdentifier.promote_from_model(default_execution_id)), execution_date=_datetime.datetime.utcnow(), - stats=_mock_stats.MockStats(), + stats=mock_stats.MockStats(), logging=_logging, tmp_dir=user_space_path, ) diff --git a/flytekit/core/data_persistence.py b/flytekit/core/data_persistence.py index b121d053fa..00c233dd8b 100644 --- a/flytekit/core/data_persistence.py +++ b/flytekit/core/data_persistence.py @@ -32,8 +32,8 @@ from typing import Dict, Union from uuid import UUID -from flytekit.common.exceptions.user import FlyteAssertion -from flytekit.common.utils import PerformanceTimer +from flytekit.core.utils import PerformanceTimer +from flytekit.exceptions.user import FlyteAssertion from flytekit.interfaces.random import random from flytekit.loggers import logger diff --git a/flytekit/common/mixins/hash.py b/flytekit/core/hash.py similarity index 100% rename from flytekit/common/mixins/hash.py rename to flytekit/core/hash.py diff --git a/flytekit/core/interface.py b/flytekit/core/interface.py index 263ac05fb7..b158a9434c 100644 --- a/flytekit/core/interface.py +++ b/flytekit/core/interface.py @@ -8,10 +8,10 @@ from collections import OrderedDict from typing import Any, Dict, Generator, List, Optional, Tuple, Type, TypeVar, Union -from flytekit.common.exceptions.user import FlyteValidationException from flytekit.core import context_manager from flytekit.core.docstring import Docstring from flytekit.core.type_engine import TypeEngine +from flytekit.exceptions.user import FlyteValidationException from flytekit.loggers import logger from flytekit.models import interface as _interface_models from flytekit.types.pickle import FlytePickle diff --git a/flytekit/core/map_task.py b/flytekit/core/map_task.py index 42645a4797..29838cffcd 100644 --- a/flytekit/core/map_task.py +++ b/flytekit/core/map_task.py @@ -8,12 +8,12 @@ from itertools import count from typing import Any, Dict, List, Optional, Type -from flytekit.common.constants import SdkTaskType -from flytekit.common.exceptions import scopes as exception_scopes from flytekit.core.base_task import PythonTask +from flytekit.core.constants import SdkTaskType from flytekit.core.context_manager import ExecutionState, FlyteContext, FlyteContextManager, SerializationSettings from flytekit.core.interface import transform_interface_to_list_interface from flytekit.core.python_function_task import PythonFunctionTask +from flytekit.exceptions import scopes as exception_scopes from flytekit.models.array_job import ArrayJob from flytekit.models.interface import Variable from flytekit.models.task import Container, K8sPod, Sql diff --git a/flytekit/engines/unit/mock_stats.py b/flytekit/core/mock_stats.py similarity index 100% rename from flytekit/engines/unit/mock_stats.py rename to flytekit/core/mock_stats.py diff --git a/flytekit/core/node.py b/flytekit/core/node.py index 0567e7b64a..cf2625b87a 100644 --- a/flytekit/core/node.py +++ b/flytekit/core/node.py @@ -4,8 +4,8 @@ import typing from typing import Any, List -from flytekit.common.utils import _dnsify from flytekit.core.resources import Resources +from flytekit.core.utils import _dnsify from flytekit.models import literals as _literal_models from flytekit.models.core import workflow as _workflow_model from flytekit.models.task import Resources as _resources_model diff --git a/flytekit/core/node_creation.py b/flytekit/core/node_creation.py index 120e9995ba..6eb4f51d81 100644 --- a/flytekit/core/node_creation.py +++ b/flytekit/core/node_creation.py @@ -3,13 +3,13 @@ import collections from typing import Type, Union -from flytekit.common.exceptions import user as _user_exceptions from flytekit.core.base_task import PythonTask from flytekit.core.context_manager import BranchEvalMode, ExecutionState, FlyteContext from flytekit.core.launch_plan import LaunchPlan from flytekit.core.node import Node from flytekit.core.promise import VoidPromise from flytekit.core.workflow import WorkflowBase +from flytekit.exceptions import user as _user_exceptions from flytekit.loggers import logger # This file exists instead of moving to node.py because it needs Task/Workflow/LaunchPlan and those depend on Node diff --git a/flytekit/core/promise.py b/flytekit/core/promise.py index 7c9b780e19..17a4837432 100644 --- a/flytekit/core/promise.py +++ b/flytekit/core/promise.py @@ -7,8 +7,7 @@ from typing_extensions import Protocol -from flytekit.common import constants as _common_constants -from flytekit.common.exceptions import user as _user_exceptions +from flytekit.core import constants as _common_constants from flytekit.core import context_manager as _flyte_context from flytekit.core import interface as flyte_interface from flytekit.core import type_engine @@ -16,6 +15,7 @@ from flytekit.core.interface import Interface from flytekit.core.node import Node from flytekit.core.type_engine import DictTransformer, ListTransformer, TypeEngine +from flytekit.exceptions import user as _user_exceptions from flytekit.models import interface as _interface_models from flytekit.models import literals as _literal_models from flytekit.models import literals as _literals_models diff --git a/flytekit/core/python_auto_container.py b/flytekit/core/python_auto_container.py index 689f5781dc..0226760f08 100644 --- a/flytekit/core/python_auto_container.py +++ b/flytekit/core/python_auto_container.py @@ -4,12 +4,12 @@ import re from typing import Callable, Dict, List, Optional, TypeVar -from flytekit.common.tasks.raw_container import _get_container_definition from flytekit.core.base_task import PythonTask, TaskResolverMixin from flytekit.core.context_manager import FlyteContextManager, ImageConfig, SerializationSettings from flytekit.core.resources import Resources, ResourceSpec from flytekit.core.tracked_abc import FlyteTrackedABC from flytekit.core.tracker import TrackedInstance +from flytekit.core.utils import _get_container_definition from flytekit.loggers import logger from flytekit.models import task as _task_model from flytekit.models.security import Secret, SecurityContext diff --git a/flytekit/core/python_customized_container_task.py b/flytekit/core/python_customized_container_task.py index eaeb509d2e..c5a716c3cb 100644 --- a/flytekit/core/python_customized_container_task.py +++ b/flytekit/core/python_customized_container_task.py @@ -5,13 +5,12 @@ from flyteidl.core import tasks_pb2 as _tasks_pb2 -from flytekit.common import utils as common_utils -from flytekit.common.tasks.raw_container import _get_container_definition from flytekit.core.base_task import PythonTask, Task, TaskResolverMixin from flytekit.core.context_manager import FlyteContext, Image, ImageConfig, SerializationSettings from flytekit.core.resources import Resources, ResourceSpec from flytekit.core.shim_task import ExecutableTemplateShimTask, ShimTaskExecutor from flytekit.core.tracker import TrackedInstance +from flytekit.core.utils import _get_container_definition, load_proto_from_file from flytekit.loggers import logger from flytekit.models import task as _task_model from flytekit.models.core import identifier as identifier_models @@ -232,7 +231,7 @@ def load_task(self, loader_args: List[str]) -> ExecutableTemplateShimTask: ctx = FlyteContext.current_context() task_template_local_path = os.path.join(ctx.execution_state.working_dir, "task_template.pb") # type: ignore ctx.file_access.get_data(loader_args[0], task_template_local_path) - task_template_proto = common_utils.load_proto_from_file(_tasks_pb2.TaskTemplate, task_template_local_path) + task_template_proto = load_proto_from_file(_tasks_pb2.TaskTemplate, task_template_local_path) task_template_model = _task_model.TaskTemplate.from_flyte_idl(task_template_proto) executor_class = load_object_from_module(loader_args[1]) diff --git a/flytekit/core/python_function_task.py b/flytekit/core/python_function_task.py index 25a363b070..fa98b7ca89 100644 --- a/flytekit/core/python_function_task.py +++ b/flytekit/core/python_function_task.py @@ -19,7 +19,6 @@ from enum import Enum from typing import Any, Callable, List, Optional, TypeVar, Union -from flytekit.common.exceptions import scopes as exception_scopes from flytekit.core.base_task import Task, TaskResolverMixin from flytekit.core.context_manager import ExecutionState, FastSerializationSettings, FlyteContext, FlyteContextManager from flytekit.core.docstring import Docstring @@ -32,6 +31,7 @@ WorkflowMetadata, WorkflowMetadataDefaults, ) +from flytekit.exceptions import scopes as exception_scopes from flytekit.loggers import logger from flytekit.models import dynamic_job as _dynamic_job from flytekit.models import literals as _literal_models @@ -178,7 +178,7 @@ def compile_into_workflow( with FlyteContextManager.with_context(ctx.with_compilation_state(cs)): # TODO: Resolve circular import - from flytekit.common.translator import get_serializable + from flytekit.tools.translator import get_serializable workflow_metadata = WorkflowMetadata(on_failure=WorkflowFailurePolicy.FAIL_IMMEDIATELY) defaults = WorkflowMetadataDefaults( diff --git a/flytekit/core/reference.py b/flytekit/core/reference.py index ed16ba6353..6a88549c43 100644 --- a/flytekit/core/reference.py +++ b/flytekit/core/reference.py @@ -2,10 +2,10 @@ from typing import Dict, Type -from flytekit.common.exceptions.user import FlyteValidationException from flytekit.core.launch_plan import ReferenceLaunchPlan from flytekit.core.task import ReferenceTask from flytekit.core.workflow import ReferenceWorkflow +from flytekit.exceptions.user import FlyteValidationException from flytekit.models.core import identifier as _identifier_model diff --git a/flytekit/core/reference_entity.py b/flytekit/core/reference_entity.py index 22090838ff..42be1313cc 100644 --- a/flytekit/core/reference_entity.py +++ b/flytekit/core/reference_entity.py @@ -2,7 +2,6 @@ from dataclasses import dataclass from typing import Any, Dict, Optional, Tuple, Type, Union -from flytekit.common.exceptions import user as _user_exceptions from flytekit.core.context_manager import BranchEvalMode, ExecutionState, FlyteContext from flytekit.core.interface import Interface, transform_interface_to_typed_interface from flytekit.core.promise import ( @@ -14,6 +13,7 @@ translate_inputs_to_literals, ) from flytekit.core.type_engine import TypeEngine +from flytekit.exceptions import user as _user_exceptions from flytekit.loggers import logger from flytekit.models import interface as _interface_models from flytekit.models import literals as _literal_models diff --git a/flytekit/core/shim_task.py b/flytekit/core/shim_task.py index 637cfe4b47..ad6d39ef38 100644 --- a/flytekit/core/shim_task.py +++ b/flytekit/core/shim_task.py @@ -2,9 +2,10 @@ from typing import Any, Generic, Type, TypeVar, Union -from flytekit import ExecutionParameters, FlyteContext, FlyteContextManager, logger +from flytekit.core.context_manager import ExecutionParameters, FlyteContext, FlyteContextManager from flytekit.core.tracker import TrackedInstance from flytekit.core.type_engine import TypeEngine +from flytekit.loggers import logger from flytekit.models import dynamic_job as _dynamic_job from flytekit.models import literals as _literal_models from flytekit.models import task as _task_model diff --git a/flytekit/core/tracker.py b/flytekit/core/tracker.py index cd2e8be02b..56f145b4b6 100644 --- a/flytekit/core/tracker.py +++ b/flytekit/core/tracker.py @@ -4,7 +4,7 @@ import logging as _logging from typing import Callable -from flytekit.common.exceptions import system as _system_exceptions +from flytekit.exceptions import system as _system_exceptions class InstanceTrackingMeta(type): diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 9c38120e69..2d41512c8e 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -26,10 +26,9 @@ from marshmallow_enum import EnumField, LoadDumpOptions from marshmallow_jsonschema import JSONSchema -from flytekit.common.exceptions import user as user_exceptions -from flytekit.common.types import primitives as _primitives from flytekit.core.context_manager import FlyteContext from flytekit.core.type_helpers import load_type_from_tag +from flytekit.exceptions import user as user_exceptions from flytekit.loggers import logger from flytekit.models import interface as _interface_models from flytekit.models import types as _type_models @@ -290,7 +289,7 @@ def get_literal_type(self, t: Type[T]) -> LiteralType: f"evaluation doesn't work with json dataclasses" ) - return _primitives.Generic.to_flyte_literal_type(metadata=schema) + return _type_models.LiteralType(simple=_type_models.SimpleType.STRUCT, metadata=schema) def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], expected: LiteralType) -> Literal: if not dataclasses.is_dataclass(python_val): @@ -923,7 +922,7 @@ def get_literal_type(self, t: Type[dict]) -> LiteralType: return _type_models.LiteralType(map_value_type=sub_type) except Exception as e: raise ValueError(f"Type of Generic List type is not supported, {e}") - return _primitives.Generic.to_flyte_literal_type() + return _type_models.LiteralType(simple=_type_models.SimpleType.STRUCT) def to_literal( self, ctx: FlyteContext, python_val: typing.Any, python_type: Type[dict], expected: LiteralType @@ -1150,7 +1149,7 @@ def _register_default_type_transformers(): SimpleTransformer( "int", int, - _primitives.Integer.to_flyte_literal_type(), + _type_models.LiteralType(simple=_type_models.SimpleType.INTEGER), lambda x: Literal(scalar=Scalar(primitive=Primitive(integer=x))), lambda x: x.scalar.primitive.integer, ) @@ -1160,7 +1159,7 @@ def _register_default_type_transformers(): SimpleTransformer( "float", float, - _primitives.Float.to_flyte_literal_type(), + _type_models.LiteralType(simple=_type_models.SimpleType.FLOAT), lambda x: Literal(scalar=Scalar(primitive=Primitive(float_value=x))), _check_and_covert_float, ) @@ -1170,7 +1169,7 @@ def _register_default_type_transformers(): SimpleTransformer( "bool", bool, - _primitives.Boolean.to_flyte_literal_type(), + _type_models.LiteralType(simple=_type_models.SimpleType.BOOLEAN), lambda x: Literal(scalar=Scalar(primitive=Primitive(boolean=x))), lambda x: x.scalar.primitive.boolean, ) @@ -1180,7 +1179,7 @@ def _register_default_type_transformers(): SimpleTransformer( "str", str, - _primitives.String.to_flyte_literal_type(), + _type_models.LiteralType(simple=_type_models.SimpleType.STRING), lambda x: Literal(scalar=Scalar(primitive=Primitive(string_value=x))), lambda x: x.scalar.primitive.string_value, ) @@ -1190,7 +1189,7 @@ def _register_default_type_transformers(): SimpleTransformer( "datetime", _datetime.datetime, - _primitives.Datetime.to_flyte_literal_type(), + _type_models.LiteralType(simple=_type_models.SimpleType.DATETIME), lambda x: Literal(scalar=Scalar(primitive=Primitive(datetime=x))), lambda x: x.scalar.primitive.datetime, ) @@ -1200,7 +1199,7 @@ def _register_default_type_transformers(): SimpleTransformer( "timedelta", _datetime.timedelta, - _primitives.Timedelta.to_flyte_literal_type(), + _type_models.LiteralType(simple=_type_models.SimpleType.DURATION), lambda x: Literal(scalar=Scalar(primitive=Primitive(duration=x))), lambda x: x.scalar.primitive.duration, ) diff --git a/flytekit/common/utils.py b/flytekit/core/utils.py similarity index 57% rename from flytekit/common/utils.py rename to flytekit/core/utils.py index 1d0395fd90..71dd45f581 100644 --- a/flytekit/common/utils.py +++ b/flytekit/core/utils.py @@ -5,10 +5,10 @@ import time as _time from hashlib import sha224 as _sha224 from pathlib import Path +from typing import Dict, List -import flytekit as _flytekit -from flytekit.configuration import sdk as _sdk_config -from flytekit.models.core import identifier as _identifier +from flytekit.configuration import resources as _resource_config +from flytekit.models import task as _task_models def _dnsify(value: str) -> str: @@ -49,6 +49,84 @@ def _dnsify(value: str) -> str: return res +def _get_container_definition( + image: str, + command: List[str], + args: List[str], + data_loading_config: _task_models.DataLoadingConfig, + storage_request: str = None, + ephemeral_storage_request: str = None, + cpu_request: str = None, + gpu_request: str = None, + memory_request: str = None, + storage_limit: str = None, + ephemeral_storage_limit: str = None, + cpu_limit: str = None, + gpu_limit: str = None, + memory_limit: str = None, + environment: Dict[str, str] = None, +) -> _task_models.Container: + storage_limit = storage_limit or _resource_config.DEFAULT_STORAGE_LIMIT.get() + storage_request = storage_request or _resource_config.DEFAULT_STORAGE_REQUEST.get() + ephemeral_storage_limit = ephemeral_storage_limit or _resource_config.DEFAULT_EPHEMERAL_STORAGE_LIMIT.get() + ephemeral_storage_request = ephemeral_storage_request or _resource_config.DEFAULT_EPHEMERAL_STORAGE_REQUEST.get() + cpu_limit = cpu_limit or _resource_config.DEFAULT_CPU_LIMIT.get() + cpu_request = cpu_request or _resource_config.DEFAULT_CPU_REQUEST.get() + gpu_limit = gpu_limit or _resource_config.DEFAULT_GPU_LIMIT.get() + gpu_request = gpu_request or _resource_config.DEFAULT_GPU_REQUEST.get() + memory_limit = memory_limit or _resource_config.DEFAULT_MEMORY_LIMIT.get() + memory_request = memory_request or _resource_config.DEFAULT_MEMORY_REQUEST.get() + + requests = [] + if storage_request: + requests.append( + _task_models.Resources.ResourceEntry(_task_models.Resources.ResourceName.STORAGE, storage_request) + ) + if ephemeral_storage_request: + requests.append( + _task_models.Resources.ResourceEntry( + _task_models.Resources.ResourceName.EPHEMERAL_STORAGE, ephemeral_storage_request + ) + ) + if cpu_request: + requests.append(_task_models.Resources.ResourceEntry(_task_models.Resources.ResourceName.CPU, cpu_request)) + if gpu_request: + requests.append(_task_models.Resources.ResourceEntry(_task_models.Resources.ResourceName.GPU, gpu_request)) + if memory_request: + requests.append( + _task_models.Resources.ResourceEntry(_task_models.Resources.ResourceName.MEMORY, memory_request) + ) + + limits = [] + if storage_limit: + limits.append(_task_models.Resources.ResourceEntry(_task_models.Resources.ResourceName.STORAGE, storage_limit)) + if ephemeral_storage_limit: + limits.append( + _task_models.Resources.ResourceEntry( + _task_models.Resources.ResourceName.EPHEMERAL_STORAGE, ephemeral_storage_limit + ) + ) + if cpu_limit: + limits.append(_task_models.Resources.ResourceEntry(_task_models.Resources.ResourceName.CPU, cpu_limit)) + if gpu_limit: + limits.append(_task_models.Resources.ResourceEntry(_task_models.Resources.ResourceName.GPU, gpu_limit)) + if memory_limit: + limits.append(_task_models.Resources.ResourceEntry(_task_models.Resources.ResourceName.MEMORY, memory_limit)) + + if environment is None: + environment = {} + + return _task_models.Container( + image=image, + command=command, + args=args, + resources=_task_models.Resources(limits=limits, requests=requests), + env=environment, + config={}, + data_loading_config=data_loading_config, + ) + + def load_proto_from_file(pb2_type, path): with open(path, "rb") as reader: out = pb2_type() @@ -62,10 +140,6 @@ def write_proto_to_file(proto, path): writer.write(proto.SerializeToString()) -def get_version_message(): - return "Welcome to Flyte! Version: {}".format(_flytekit.__version__) - - class Directory(object): def __init__(self, path): """ @@ -160,65 +234,3 @@ def __exit__(self, exc_type, exc_val, exc_tb): end_process_time - self._start_process_time, ) ) - - -class ExitStack(object): - def __init__(self, entered_stack=None): - self._contexts = entered_stack - - def enter_context(self, context): - out = context.__enter__() - self._contexts.append(context) - return out - - def pop_all(self): - entered_stack = self._contexts - self._contexts = None - return ExitStack(entered_stack=entered_stack) - - def __enter__(self): - if self._contexts is not None: - raise Exception("A non-empty context stack cannot be entered.") - self._contexts = [] - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - first_exception = None - if self._contexts is not None: - while len(self._contexts) > 0: - try: - self._contexts.pop().__exit__(exc_type, exc_val, exc_tb) - except Exception as ex: - # Catch all to try to clean up all exits before re-raising the first exception - if first_exception is None: - first_exception = ex - if first_exception is not None: - raise first_exception - return False - - -def fqdn(module, name, entity_type=None): - """ - :param Text module: - :param Text name: - :param int entity_type: _identifier.ResourceType enum - :rtype: Text - """ - fmt = _sdk_config.NAME_FORMAT.get() - if entity_type == _identifier.ResourceType.WORKFLOW: - fmt = _sdk_config.WORKFLOW_NAME_FORMAT.get() or fmt - elif entity_type == _identifier.ResourceType.TASK: - fmt = _sdk_config.TASK_NAME_FORMAT.get() or fmt - elif entity_type == _identifier.ResourceType.LAUNCH_PLAN: - fmt = _sdk_config.LAUNCH_PLAN_NAME_FORMAT.get() or fmt - return fmt.format(module=module, name=name) - - -def fqdn_safe(module, key, entity_type=None): - """ - :param Text module: - :param Text key: - :param int entity_type: _identifier.ResourceType enum - :rtype: Text - """ - return _dnsify(fqdn(module, key, entity_type=entity_type)) diff --git a/flytekit/core/workflow.py b/flytekit/core/workflow.py index ffa6aae934..77d7a6a936 100644 --- a/flytekit/core/workflow.py +++ b/flytekit/core/workflow.py @@ -5,9 +5,7 @@ from functools import update_wrapper from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union -from flytekit.common import constants as _common_constants -from flytekit.common.exceptions import scopes as exception_scopes -from flytekit.common.exceptions.user import FlyteValidationException, FlyteValueException +from flytekit.core import constants as _common_constants from flytekit.core.base_task import PythonTask from flytekit.core.class_based_resolver import ClassStorageTaskResolver from flytekit.core.condition import ConditionalSection @@ -34,6 +32,8 @@ from flytekit.core.python_auto_container import PythonAutoContainerTask from flytekit.core.reference_entity import ReferenceEntity, WorkflowReference from flytekit.core.type_engine import TypeEngine +from flytekit.exceptions import scopes as exception_scopes +from flytekit.exceptions.user import FlyteValidationException, FlyteValueException from flytekit.loggers import logger from flytekit.models import interface as _interface_models from flytekit.models import literals as _literal_models diff --git a/flytekit/engines/__init__.py b/flytekit/engines/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/flytekit/engines/common.py b/flytekit/engines/common.py deleted file mode 100644 index ed1dfd9d37..0000000000 --- a/flytekit/engines/common.py +++ /dev/null @@ -1,422 +0,0 @@ -import abc as _abc - -from flytekit.models import common as _common_models - - -class BaseWorkflowExecutor(object, metaclass=_common_models.FlyteABCMeta): - """ - This class must be implemented for any engine to create, interact with, and execute workflows using the - FlyteKit SDK. - """ - - def __init__(self, sdk_workflow): - """ - :param flytekit.common.workflow.SdkWorkflow sdk_workflow: - """ - self._sdk_workflow = sdk_workflow - - @property - def sdk_workflow(self): - """ - :rtype: flytekit.common.workflow.SdkWorkflow - """ - return self._sdk_workflow - - @_abc.abstractmethod - def register(self, identifier): - """ - Registers the workflow - :param flytekit.models.core.identifier.Identifier identifier: - """ - pass - - -class BaseWorkflowExecution(object, metaclass=_common_models.FlyteABCMeta): - """ - This class must be implemented for any engine to track and interact with the executions of workflows. - """ - - def __init__(self, sdk_wf_exec): - """ - :param flytekit.common.workflow_execution.SdkWorkflowExecution sdk_wf_exec: - """ - self._sdk_wf_exec = sdk_wf_exec - - @property - def sdk_workflow_execution(self): - """ - :rtype: flytekit.common.workflow_execution.SdkWorkflowExecution - """ - return self._sdk_wf_exec - - @_abc.abstractmethod - def get_node_executions(self, filters=None): - """ - :param list[flytekit.models.filters.Filter] filters: - :rtype: dict[Text, flytekit.common.nodes.SdkNodeExecution] - """ - pass - - @_abc.abstractmethod - def sync(self): - """ - :rtype: None - """ - pass - - @_abc.abstractmethod - def get_inputs(self): - """ - :rtype: flytekit.models.literals.LiteralMap - """ - pass - - @_abc.abstractmethod - def get_outputs(self): - """ - :rtype: flytekit.models.literals.LiteralMap - """ - pass - - @_abc.abstractmethod - def terminate(self, cause): - """ - :param Text cause: - """ - pass - - -class BaseNodeExecution(object, metaclass=_common_models.FlyteABCMeta): - def __init__(self, node_execution): - """ - :param flytekit.common.nodes.SdkNodeExecution node_execution: - """ - self._sdk_node_execution = node_execution - - @property - def sdk_node_execution(self): - """ - :rtype: flytekit.common.nodes.SdkNodeExecution - """ - return self._sdk_node_execution - - @_abc.abstractmethod - def get_task_executions(self): - """ - :rtype: list[flytekit.common.tasks.executions.SdkTaskExecution] - """ - pass - - @_abc.abstractmethod - def get_subworkflow_executions(self): - """ - :rtype: list[flytekit.common.workflow_execution.SdkWorkflowExecution] - """ - pass - - @_abc.abstractmethod - def get_inputs(self): - """ - :rtype: flytekit.models.literals.LiteralMap - """ - pass - - @_abc.abstractmethod - def get_outputs(self): - """ - :rtype: flytekit.models.literals.LiteralMap - """ - pass - - @_abc.abstractmethod - def sync(self): - """ - :rtype: None - """ - pass - - -class BaseTaskExecution(object, metaclass=_common_models.FlyteABCMeta): - def __init__(self, task_exec): - """ - :param flytekit.common.tasks.executions.SdkTaskExecution task_exec: - """ - self._sdk_task_execution = task_exec - - @property - def sdk_task_execution(self): - """ - :rtype: flytekit.common.tasks.executions.SdkTaskExecution - """ - return self._sdk_task_execution - - @_abc.abstractmethod - def get_inputs(self): - """ - :rtype: flytekit.models.literals.LiteralMap - """ - pass - - @_abc.abstractmethod - def get_outputs(self): - """ - :rtype: flytekit.models.literals.LiteralMap - """ - pass - - @_abc.abstractmethod - def sync(self): - """ - :rtype: None - """ - pass - - @_abc.abstractmethod - def get_child_executions(self, filters=None): - """ - :param list[flytekit.models.filters.Filter] filters: - :rtype: dict[Text, flytekit.common.nodes.SdkNodeExecution] - """ - pass - - -class BaseLaunchPlanLauncher(object, metaclass=_common_models.FlyteABCMeta): - def __init__(self, sdk_launch_plan): - """ - :param flytekit.common.launch_plan.SdkLaunchPlan sdk_launch_plan: - """ - self._sdk_launch_plan = sdk_launch_plan - - @property - def sdk_launch_plan(self): - """ - :rtype: flytekit.common.launch_plan.SdkLaunchPlan - """ - return self._sdk_launch_plan - - @_abc.abstractmethod - def register(self, identifier): - """ - Registers the launch plan - :param flytekit.models.core.identifier.Identifier identifier: - """ - pass - - @_abc.abstractmethod - def launch( - self, - project, - domain, - name, - inputs, - notification_overrides=None, - label_overrides=None, - annotation_overrides=None, - ): - """ - Registers the launch plan and returns the identifier. - :param Text project: - :param Text domain: - :param Text name: - :param flytekit.models.literals.LiteralMap inputs: The inputs to pass - :param list[flytekit.models.common.Notification] notification_overrides: If specified, override the - notifications. - :param flytekit.models.common.Labels label_overrides: - :param flytekit.models.common.Annotations annotation_overrides: - :rtype: flytekit.models.execution.Execution - """ - pass - - @_abc.abstractmethod - def update(self, identifier, state): - """ - :param flytekit.models.core.identifier.Identifier identifier: ID for launch plan to update - :param int state: Enum value from flytekit.models.launch_plan.LaunchPlanState - """ - pass - - -class BaseTaskExecutor(object, metaclass=_common_models.FlyteABCMeta): - def __init__(self, sdk_task): - """ - :param flytekit.common.tasks.task.SdkTask sdk_task: - """ - self._sdk_task = sdk_task - - @property - def sdk_task(self): - """ - :rtype: flytekit.common.tasks.sdk_runnable.SdkRunnableTask - """ - return self._sdk_task - - @_abc.abstractmethod - def execute(self, inputs, context=None): - """ - :param flytekit.models.literals.LiteralMap inputs: Inputs to pass to the workflow. - """ - pass - - @_abc.abstractmethod - def register(self, identifier): - """ - Registers the task - :param flytekit.models.core.identifier.Identifier identifier: - """ - pass - - @_abc.abstractmethod - def launch( - self, - project, - domain, - name=None, - inputs=None, - notification_overrides=None, - label_overrides=None, - annotation_overrides=None, - auth_role=None, - ): - """ - Executes the task as a single task execution and returns the identifier. - :param Text project: - :param Text domain: - :param Text name: - :param flytekit.models.literals.LiteralMap inputs: The inputs to pass - :param list[flytekit.models.common.Notification] notification_overrides: If specified, override the - notifications. - :param flytekit.models.common.Labels label_overrides: - :param flytekit.models.common.Annotations annotation_overrides: - :param flytekit.models.common.AuthRole auth_role: - :rtype: flytekit.models.execution.Execution - """ - pass - - -class BaseExecutionEngineFactory(object, metaclass=_common_models.FlyteABCMeta): - """ - This object should be implemented to satisfy the basic engine interface. - """ - - @_abc.abstractmethod - def get_task(self, sdk_task): - """ - :param flytekit.common.tasks.task.SdkTask sdk_task: - :rtype: BaseTaskExecutor - """ - pass - - @_abc.abstractmethod - def get_launch_plan(self, sdk_launch_plan): - """ - :param flytekit.common.launch_plan.SdkLaunchPlan sdk_launch_plan: - :rtype: BaseLaunchPlanLauncher - """ - pass - - @_abc.abstractmethod - def get_task_execution(self, task_exec): - """ - :param flytekit.common.tasks.executions.SdkTaskExecution task_exec: - :rtype: BaseTaskExecution - """ - pass - - @_abc.abstractmethod - def get_node_execution(self, node_exec): - """ - :param flytekit.common.nodes.SdkNodeExecution node_exec: - :rtype: BaseNodeExecution - """ - pass - - @_abc.abstractmethod - def get_workflow_execution(self, wf_exec): - """ - :param flytekit.common.workflow_execution.SdkWorkflowExecution wf_exec: - :rtype: BaseWorkflowExecution - """ - pass - - @_abc.abstractmethod - def fetch_workflow_execution(self, wf_exec_id): - """ - :param flytekit.models.core.identifier.WorkflowExecutionIdentifier wf_exec_id: - :rtype: flytekit.models.execution.Execution - """ - pass - - @_abc.abstractmethod - def fetch_task(self, task_id): - """ - :param flytekit.models.core.identifier.Identifier task_id: This identifier should have a resource type of kind - Task. - :rtype: flytekit.models.task.Task - """ - pass - - @_abc.abstractmethod - def fetch_latest_task(self, named_task): - """ - Fetches the latest task - :param flytekit.models.common.NamedEntityIdentifier named_task: NamedEntityIdentifier to fetch - :rtype: flytekit.models.task.Task - """ - pass - - -class EngineContext(object): - def __init__( - self, - execution_date, - tmp_dir, - stats, - execution_id, - logging, - raw_output_data_prefix=None, - ): - self._stats = stats - self._execution_date = execution_date - self._working_directory = tmp_dir - self._execution_id = execution_id - self._logging = logging - self._raw_output_data_prefix = raw_output_data_prefix - - @property - def stats(self): - """ - :rtype: flytekit.interfaces.stats.taggable.TaggableStats - """ - return self._stats - - @property - def logging(self): - """ - :rtype: TODO - """ - return self._logging - - @property - def working_directory(self): - """ - :rtype: flytekit.common.utils.AutoDeletingTempDir - """ - return self._working_directory - - @property - def execution_date(self): - """ - :rtype: datetime.datetime - """ - return self._execution_date - - @property - def execution_id(self): - """ - :rtype: flytekit.models.core.identifier.WorkflowExecutionIdentifier - """ - return self._execution_id - - @property - def raw_output_data_prefix(self) -> str: - return self._raw_output_data_prefix diff --git a/flytekit/engines/flyte/__init__.py b/flytekit/engines/flyte/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/flytekit/engines/flyte/engine.py b/flytekit/engines/flyte/engine.py deleted file mode 100644 index 7af35e20b9..0000000000 --- a/flytekit/engines/flyte/engine.py +++ /dev/null @@ -1,727 +0,0 @@ -import logging as _logging -import os as _os -import traceback as _traceback -from datetime import datetime as _datetime - -import six as _six -from deprecated import deprecated as _deprecated -from flyteidl.core import literals_pb2 as _literals_pb2 - -import flytekit -from flytekit.clients.friendly import SynchronousFlyteClient as _SynchronousFlyteClient -from flytekit.clients.helpers import iterate_node_executions as _iterate_node_executions -from flytekit.clients.helpers import iterate_task_executions as _iterate_task_executions -from flytekit.common import constants as _constants -from flytekit.common import utils as _common_utils -from flytekit.common.exceptions import scopes as _exception_scopes -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.configuration import auth as _auth_config -from flytekit.configuration import internal as _internal_config -from flytekit.configuration import platform as _platform_config -from flytekit.configuration import sdk as _sdk_config -from flytekit.engines import common as _common_engine -from flytekit.interfaces.data import data_proxy as _data_proxy -from flytekit.interfaces.stats.taggable import get_stats as _get_stats -from flytekit.models import common as _common_models -from flytekit.models import execution as _execution_models -from flytekit.models import literals as _literals -from flytekit.models import task as _task_models -from flytekit.models.admin import common as _common -from flytekit.models.admin import workflow as _workflow_model -from flytekit.models.core import errors as _error_models -from flytekit.models.core import identifier as _identifier - - -class _FlyteClientManager(object): - _CLIENT = None - - def __init__(self, *args, **kwargs): - # TODO: React to changing configs. For now this is frozen for the lifetime of the process, which covers most - # TODO: use cases. - if type(self)._CLIENT is None: - c = _SynchronousFlyteClient(*args, **kwargs) - type(self)._CLIENT = c - - @property - def client(self): - """ - :rtype: flytekit.clients.friendly.SynchronousFlyteClient - """ - return type(self)._CLIENT - - -# This is a simple helper function that ties the client together with the configuration construct. -# This will be refactored away when we move to a heavier context object. -def get_client() -> _SynchronousFlyteClient: - return _FlyteClientManager(_platform_config.URL.get(), insecure=_platform_config.INSECURE.get()).client - - -class FlyteEngineFactory(_common_engine.BaseExecutionEngineFactory): - def get_workflow(self, sdk_workflow): - """ - :param flytekit.common.workflow.SdkWorkflow sdk_workflow: - :rtype: FlyteWorkflow - """ - return FlyteWorkflow(sdk_workflow) - - def get_task(self, sdk_task): - """ - :param flytekit.common.tasks.task.SdkTask sdk_task: - :rtype: FlyteTask - """ - return FlyteTask(sdk_task) - - def get_launch_plan(self, sdk_launch_plan): - """ - :param flytekit.common.launch_plan.SdkLaunchPlan sdk_launch_plan: - :rtype: FlyteLaunchPlan - """ - return FlyteLaunchPlan(sdk_launch_plan) - - def get_task_execution(self, task_exec): - """ - :param flytekit.common.tasks.executions.SdkTaskExecution task_exec: - :rtype: FlyteTaskExecution - """ - return FlyteTaskExecution(task_exec) - - def get_node_execution(self, node_exec): - """ - :param flytekit.common.nodes.SdkNodeExecution node_exec: - :rtype: FlyteNodeExecution - """ - return FlyteNodeExecution(node_exec) - - def get_workflow_execution(self, wf_exec): - """ - :param flytekit.common.workflow_execution.SdkWorkflowExecution wf_exec: - :rtype: FlyteWorkflowExecution - """ - return FlyteWorkflowExecution(wf_exec) - - @_deprecated( - reason="Objects should access client directly, will be removed by 1.0", - version="0.13.0", - ) - def fetch_workflow_execution(self, wf_exec_id): - """ - :param flytekit.models.core.identifier.WorkflowExecutionIdentifier wf_exec_id: - :rtype: flytekit.models.execution.Execution - """ - return _FlyteClientManager( - _platform_config.URL.get(), insecure=_platform_config.INSECURE.get() - ).client.get_execution(wf_exec_id) - - @_deprecated( - reason="Objects should access client directly, will be removed by 1.0", - version="0.13.0", - ) - def fetch_task(self, task_id): - """ - Queries Admin for an existing Admin task - :param flytekit.models.core.identifier.Identifier task_id: - :rtype: flytekit.models.task.Task - """ - return _FlyteClientManager( - _platform_config.URL.get(), insecure=_platform_config.INSECURE.get() - ).client.get_task(task_id) - - @_deprecated( - reason="Objects should access client directly, will be removed by 1.0", - version="0.13.0", - ) - def fetch_latest_task(self, named_task): - """ - Fetches the latest task - :param flytekit.models.common.NamedEntityIdentifier named_task: NamedEntityIdentifier to fetch - :rtype: flytekit.models.task.Task - """ - task_list, _ = _FlyteClientManager( - _platform_config.URL.get(), insecure=_platform_config.INSECURE.get() - ).client.list_tasks_paginated( - named_task, - limit=1, - sort_by=_common.Sort("created_at", _common.Sort.Direction.DESCENDING), - ) - return task_list[0] if task_list else None - - @_deprecated( - reason="Objects should access client directly, will be removed by 1.0", - version="0.13.0", - ) - def fetch_launch_plan(self, launch_plan_id): - """ - :param flytekit.models.core.identifier.Identifier launch_plan_id: This identifier should have a resource - type of kind LaunchPlan. - :rtype: flytekit.models.launch_plan.LaunchPlan - """ - if launch_plan_id.version: - return _FlyteClientManager( - _platform_config.URL.get(), insecure=_platform_config.INSECURE.get() - ).client.get_launch_plan(launch_plan_id) - else: - named_entity_id = _common_models.NamedEntityIdentifier( - launch_plan_id.project, launch_plan_id.domain, launch_plan_id.name - ) - return _FlyteClientManager( - _platform_config.URL.get(), insecure=_platform_config.INSECURE.get() - ).client.get_active_launch_plan(named_entity_id) - - @_deprecated( - reason="Objects should access client directly, will be removed by 1.0", - version="0.13.0", - ) - def fetch_workflow(self, workflow_id): - """ - :param flytekit.models.core.identifier.Identifier workflow_id: This identifier should have a resource - type of kind LaunchPlan. - :rtype: flytekit.models.admin.workflow.Workflow - """ - return _FlyteClientManager( - _platform_config.URL.get(), insecure=_platform_config.INSECURE.get() - ).client.get_workflow(workflow_id) - - -class FlyteLaunchPlan(_common_engine.BaseLaunchPlanLauncher): - @_deprecated( - reason="Objects should access client directly, will be removed by 1.0", - version="0.13.0", - ) - def register(self, identifier): - client = _FlyteClientManager(_platform_config.URL.get(), insecure=_platform_config.INSECURE.get()).client - try: - client.create_launch_plan(identifier, self.sdk_launch_plan) - except _user_exceptions.FlyteEntityAlreadyExistsException: - pass - - @_deprecated(reason="Use launch instead", version="0.9.0") - def execute( - self, - project, - domain, - name, - inputs, - notification_overrides=None, - label_overrides=None, - annotation_overrides=None, - ): - """ - Deprecated. Use launch instead. - """ - return self.launch( - project, - domain, - name, - inputs, - notification_overrides, - label_overrides, - annotation_overrides, - ) - - @_deprecated( - reason="Objects should access client directly, will be removed by 1.0", - version="0.13.0", - ) - def launch( - self, - project, - domain, - name, - inputs, - notification_overrides=None, - label_overrides=None, - annotation_overrides=None, - ): - """ - Creates a workflow execution using parameters specified in the launch plan. - :param Text project: - :param Text domain: - :param Text name: - :param flytekit.models.literals.LiteralMap inputs: - :param list[flytekit.models.common.Notification] notification_overrides: If specified, override the - notifications. - :param flytekit.models.common.Labels label_overrides: - :param flytekit.models.common.Annotations annotation_overrides: - :rtype: flytekit.models.execution.Execution - """ - disable_all = notification_overrides == [] - if disable_all: - notification_overrides = None - else: - notification_overrides = _execution_models.NotificationList(notification_overrides or []) - disable_all = None - - try: - client = _FlyteClientManager(_platform_config.URL.get(), insecure=_platform_config.INSECURE.get()).client - exec_id = client.create_execution( - project, - domain, - name, - _execution_models.ExecutionSpec( - self.sdk_launch_plan.id, - _execution_models.ExecutionMetadata( - _execution_models.ExecutionMetadata.ExecutionMode.MANUAL, - "sdk", # TODO: get principle - 0, # TODO: Detect nesting - ), - notifications=notification_overrides, - disable_all=disable_all, - labels=label_overrides, - annotations=annotation_overrides, - ), - inputs, - ) - except _user_exceptions.FlyteEntityAlreadyExistsException: - exec_id = _identifier.WorkflowExecutionIdentifier(project, domain, name) - return client.get_execution(exec_id) - - @_deprecated( - reason="Objects should access client directly, will be removed by 1.0", - version="0.13.0", - ) - def update(self, identifier, state): - """ - :param flytekit.models.core.identifier.Identifier identifier: Identifier for launch plan to update - :param int state: Enum value from flytekit.models.launch_plan.LaunchPlanState - """ - return _FlyteClientManager( - _platform_config.URL.get(), insecure=_platform_config.INSECURE.get() - ).client.update_launch_plan(identifier, state) - - -class FlyteWorkflow(_common_engine.BaseWorkflowExecutor): - @_deprecated( - reason="Objects should access client directly, will be removed by 1.0", - version="0.13.0", - ) - def register(self, identifier): - client = _FlyteClientManager(_platform_config.URL.get(), insecure=_platform_config.INSECURE.get()).client - try: - sub_workflows = self.sdk_workflow.get_sub_workflows() - return client.create_workflow( - identifier, - _workflow_model.WorkflowSpec( - self.sdk_workflow, - sub_workflows, - ), - ) - except _user_exceptions.FlyteEntityAlreadyExistsException: - pass - - -class FlyteTask(_common_engine.BaseTaskExecutor): - @_deprecated( - reason="Objects should access client directly, will be removed by 1.0", - version="0.13.0", - ) - def register(self, identifier): - client = _FlyteClientManager(_platform_config.URL.get(), insecure=_platform_config.INSECURE.get()).client - try: - client.create_task(identifier, _task_models.TaskSpec(self.sdk_task)) - except _user_exceptions.FlyteEntityAlreadyExistsException: - pass - - def execute(self, inputs, context=None): - """ - Just execute the task and write the outputs to where they belong - :param flytekit.models.literals.LiteralMap inputs: - :param dict[Text, Text] context: - :rtype: dict[Text, flytekit.models.common.FlyteIdlEntity] - """ - with _common_utils.AutoDeletingTempDir("engine_dir") as temp_dir: - with _common_utils.AutoDeletingTempDir("task_dir") as task_dir: - with _data_proxy.LocalWorkingDirectoryContext(task_dir): - raw_output_data_prefix = context.get("raw_output_data_prefix", None) - with _data_proxy.RemoteDataContext(raw_output_data_prefix_override=raw_output_data_prefix): - output_file_dict = dict() - - # This sets the logging level for user code and is the only place an sdk setting gets - # used at runtime. Optionally, Propeller can set an internal config setting which - # takes precedence. - log_level = _internal_config.LOGGING_LEVEL.get() or _sdk_config.LOGGING_LEVEL.get() - _logging.getLogger().setLevel(log_level) - - try: - output_file_dict = self.sdk_task.execute( - _common_engine.EngineContext( - execution_id=_identifier.WorkflowExecutionIdentifier( - project=_internal_config.EXECUTION_PROJECT.get(), - domain=_internal_config.EXECUTION_DOMAIN.get(), - name=_internal_config.EXECUTION_NAME.get(), - ), - execution_date=_datetime.utcnow(), - stats=_get_stats( - # Stats metric path will be: - # registration_project.registration_domain.app.module.task_name.user_stats - # and it will be tagged with execution-level values for project/domain/wf/lp - "{}.{}.{}.user_stats".format( - _internal_config.TASK_PROJECT.get() or _internal_config.PROJECT.get(), - _internal_config.TASK_DOMAIN.get() or _internal_config.DOMAIN.get(), - _internal_config.TASK_NAME.get() or _internal_config.NAME.get(), - ), - tags={ - "exec_project": _internal_config.EXECUTION_PROJECT.get(), - "exec_domain": _internal_config.EXECUTION_DOMAIN.get(), - "exec_workflow": _internal_config.EXECUTION_WORKFLOW.get(), - "exec_launchplan": _internal_config.EXECUTION_LAUNCHPLAN.get(), - "api_version": flytekit.__version__, - }, - ), - logging=_logging, - tmp_dir=task_dir, - raw_output_data_prefix=context["raw_output_data_prefix"] - if "raw_output_data_prefix" in context - else None, - ), - inputs, - ) - except _exception_scopes.FlyteScopedException as e: - _logging.error("!!! Begin Error Captured by Flyte !!!") - output_file_dict[_constants.ERROR_FILE_NAME] = _error_models.ErrorDocument( - _error_models.ContainerError(e.error_code, e.verbose_message, e.kind, 0) - ) - _logging.error(e.verbose_message) - _logging.error("!!! End Error Captured by Flyte !!!") - except Exception: - _logging.error("!!! Begin Unknown System Error Captured by Flyte !!!") - exc_str = _traceback.format_exc() - output_file_dict[_constants.ERROR_FILE_NAME] = _error_models.ErrorDocument( - _error_models.ContainerError( - "SYSTEM:Unknown", exc_str, _error_models.ContainerError.Kind.RECOVERABLE, 0 - ) - ) - _logging.error(exc_str) - _logging.error("!!! End Error Captured by Flyte !!!") - finally: - for k, v in _six.iteritems(output_file_dict): - _common_utils.write_proto_to_file(v.to_flyte_idl(), _os.path.join(temp_dir.name, k)) - _data_proxy.Data.put_data( - temp_dir.name, - context["output_prefix"], - is_multipart=True, - ) - - @_deprecated( - reason="Objects should access client directly, will be removed by 1.0", - version="0.13.0", - ) - def launch( - self, - project, - domain, - name=None, - inputs=None, - notification_overrides=None, - label_overrides=None, - annotation_overrides=None, - auth_role=None, - ): - """ - Executes the task as a single task execution and returns the identifier. - :param Text project: - :param Text domain: - :param Text name: - :param flytekit.models.literals.LiteralMap inputs: The inputs to pass - :param list[flytekit.models.common.Notification] notification_overrides: If specified, override the - notifications. - :param flytekit.models.common.Labels label_overrides: - :param flytekit.models.common.Annotations annotation_overrides: - :param flytekit.models.common.AuthRole auth_role: - :rtype: flytekit.models.execution.Execution - """ - disable_all = notification_overrides == [] - if disable_all: - notification_overrides = None - else: - notification_overrides = _execution_models.NotificationList(notification_overrides or []) - disable_all = None - - if not auth_role: - assumable_iam_role = _auth_config.ASSUMABLE_IAM_ROLE.get() - kubernetes_service_account = _auth_config.KUBERNETES_SERVICE_ACCOUNT.get() - - if not (assumable_iam_role or kubernetes_service_account): - _logging.warning( - "Using deprecated `role` from config. " - "Please update your config to use `assumable_iam_role` instead" - ) - assumable_iam_role = _sdk_config.ROLE.get() - auth_role = _common_models.AuthRole( - assumable_iam_role=assumable_iam_role, - kubernetes_service_account=kubernetes_service_account, - ) - - try: - # TODO(katrogan): Add handling to register the underlying task if it's not already. - client = _FlyteClientManager(_platform_config.URL.get(), insecure=_platform_config.INSECURE.get()).client - exec_id = client.create_execution( - project, - domain, - name, - _execution_models.ExecutionSpec( - self.sdk_task.id, - _execution_models.ExecutionMetadata( - _execution_models.ExecutionMetadata.ExecutionMode.MANUAL, - "sdk", # TODO: get principle - 0, # TODO: Detect nesting - ), - notifications=notification_overrides, - disable_all=disable_all, - labels=label_overrides, - annotations=annotation_overrides, - auth_role=auth_role, - ), - inputs, - ) - except _user_exceptions.FlyteEntityAlreadyExistsException: - exec_id = _identifier.WorkflowExecutionIdentifier(project, domain, name) - return client.get_execution(exec_id) - - -class FlyteWorkflowExecution(_common_engine.BaseWorkflowExecution): - @_deprecated( - reason="Objects should access client directly, will be removed by 1.0", - version="0.13.0", - ) - def get_node_executions(self, filters=None): - """ - :param list[flytekit.models.filters.Filter] filters: - :rtype: dict[Text, flytekit.common.nodes.SdkNodeExecution] - """ - client = _FlyteClientManager(_platform_config.URL.get(), insecure=_platform_config.INSECURE.get()).client - return { - v.id.node_id: v for v in _iterate_node_executions(client, self.sdk_workflow_execution.id, filters=filters) - } - - @_deprecated( - reason="Objects should access client directly, will be removed by 1.0", - version="0.13.0", - ) - def sync(self): - """ - :rtype: None - """ - client = _FlyteClientManager(_platform_config.URL.get(), insecure=_platform_config.INSECURE.get()).client - self.sdk_workflow_execution._closure = client.get_execution(self.sdk_workflow_execution.id).closure - - @_deprecated( - reason="Objects should access client directly, will be removed by 1.0", - version="0.13.0", - ) - def get_inputs(self): - """ - :rtype: flytekit.models.literals.LiteralMap - """ - client = _FlyteClientManager(_platform_config.URL.get(), insecure=_platform_config.INSECURE.get()).client - execution_data = client.get_execution_data(self.sdk_workflow_execution.id) - - # Inputs are returned inline unless they are too big, in which case a url blob pointing to them is returned. - if bool(execution_data.full_inputs.literals): - return execution_data.full_inputs - - if execution_data.inputs.bytes > 0: - with _common_utils.AutoDeletingTempDir() as t: - tmp_name = _os.path.join(t.name, "inputs.pb") - _data_proxy.Data.get_data(execution_data.inputs.url, tmp_name) - return _literals.LiteralMap.from_flyte_idl( - _common_utils.load_proto_from_file(_literals_pb2.LiteralMap, tmp_name) - ) - return _literals.LiteralMap({}) - - @_deprecated( - reason="Objects should access client directly, will be removed by 1.0", - version="0.13.0", - ) - def get_outputs(self): - """ - :rtype: flytekit.models.literals.LiteralMap - """ - client = _FlyteClientManager(_platform_config.URL.get(), insecure=_platform_config.INSECURE.get()).client - execution_data = client.get_execution_data(self.sdk_workflow_execution.id) - - # Outputs are returned inline unless they are too big, in which case a url blob pointing to them is returned. - if bool(execution_data.full_outputs.literals): - return execution_data.full_outputs - - if execution_data.outputs.bytes > 0: - with _common_utils.AutoDeletingTempDir() as t: - tmp_name = _os.path.join(t.name, "outputs.pb") - _data_proxy.Data.get_data(execution_data.outputs.url, tmp_name) - return _literals.LiteralMap.from_flyte_idl( - _common_utils.load_proto_from_file(_literals_pb2.LiteralMap, tmp_name) - ) - return _literals.LiteralMap({}) - - @_deprecated( - reason="Objects should access client directly, will be removed by 1.0", - version="0.13.0", - ) - def terminate(self, cause): - """ - :param Text cause: - """ - _FlyteClientManager( - _platform_config.URL.get(), insecure=_platform_config.INSECURE.get() - ).client.terminate_execution(self.sdk_workflow_execution.id, cause) - - -class FlyteNodeExecution(_common_engine.BaseNodeExecution): - @_deprecated( - reason="Objects should access client directly, will be removed by 1.0", - version="0.13.0", - ) - def get_task_executions(self): - """ - :rtype: list[flytekit.common.tasks.executions.SdkTaskExecution] - """ - client = _FlyteClientManager(_platform_config.URL.get(), insecure=_platform_config.INSECURE.get()).client - return list(_iterate_task_executions(client, self.sdk_node_execution.id)) - - @_deprecated( - reason="Objects should access client directly, will be removed by 1.0", - version="0.13.0", - ) - def get_subworkflow_executions(self): - """ - :rtype: list[flytekit.common.workflow_execution.SdkWorkflowExecution] - """ - raise NotImplementedError("Cannot retrieve sub-workflow information from a node execution yet.") - - @_deprecated( - reason="Objects should access client directly, will be removed by 1.0", - version="0.13.0", - ) - def get_inputs(self): - """ - :rtype: flytekit.models.literals.LiteralMap - """ - client = _FlyteClientManager(_platform_config.URL.get(), insecure=_platform_config.INSECURE.get()).client - execution_data = client.get_node_execution_data(self.sdk_node_execution.id) - - # Inputs are returned inline unless they are too big, in which case a url blob pointing to them is returned. - if bool(execution_data.full_inputs.literals): - return execution_data.full_inputs - - if execution_data.inputs.bytes > 0: - with _common_utils.AutoDeletingTempDir() as t: - tmp_name = _os.path.join(t.name, "inputs.pb") - _data_proxy.Data.get_data(execution_data.inputs.url, tmp_name) - return _literals.LiteralMap.from_flyte_idl( - _common_utils.load_proto_from_file(_literals_pb2.LiteralMap, tmp_name) - ) - return _literals.LiteralMap({}) - - @_deprecated( - reason="Objects should access client directly, will be removed by 1.0", - version="0.13.0", - ) - def get_outputs(self): - """ - :rtype: flytekit.models.literals.LiteralMap - """ - client = _FlyteClientManager(_platform_config.URL.get(), insecure=_platform_config.INSECURE.get()).client - execution_data = client.get_node_execution_data(self.sdk_node_execution.id) - - # Outputs are returned inline unless they are too big, in which case a url blob pointing to them is returned. - if bool(execution_data.full_outputs.literals): - return execution_data.full_outputs - - if execution_data.outputs.bytes > 0: - with _common_utils.AutoDeletingTempDir() as t: - tmp_name = _os.path.join(t.name, "outputs.pb") - _data_proxy.Data.get_data(execution_data.outputs.url, tmp_name) - return _literals.LiteralMap.from_flyte_idl( - _common_utils.load_proto_from_file(_literals_pb2.LiteralMap, tmp_name) - ) - return _literals.LiteralMap({}) - - @_deprecated( - reason="Objects should access client directly, will be removed by 1.0", - version="0.13.0", - ) - def sync(self): - """ - :rtype: None - """ - client = _FlyteClientManager(_platform_config.URL.get(), insecure=_platform_config.INSECURE.get()).client - self.sdk_node_execution._closure = client.get_node_execution(self.sdk_node_execution.id).closure - - -class FlyteTaskExecution(_common_engine.BaseTaskExecution): - @_deprecated( - reason="Objects should access client directly, will be removed by 1.0", - version="0.13.0", - ) - def get_inputs(self): - """ - :rtype: flytekit.models.literals.LiteralMap - """ - client = _FlyteClientManager(_platform_config.URL.get(), insecure=_platform_config.INSECURE.get()).client - execution_data = client.get_task_execution_data(self.sdk_task_execution.id) - - # Inputs are returned inline unless they are too big, in which case a url blob pointing to them is returned. - if bool(execution_data.full_inputs.literals): - return execution_data.full_inputs - - if execution_data.inputs.bytes > 0: - with _common_utils.AutoDeletingTempDir() as t: - tmp_name = _os.path.join(t.name, "inputs.pb") - _data_proxy.Data.get_data(execution_data.inputs.url, tmp_name) - return _literals.LiteralMap.from_flyte_idl( - _common_utils.load_proto_from_file(_literals_pb2.LiteralMap, tmp_name) - ) - return _literals.LiteralMap({}) - - @_deprecated( - reason="Objects should access client directly, will be removed by 1.0", - version="0.13.0", - ) - def get_outputs(self): - """ - :rtype: flytekit.models.literals.LiteralMap - """ - client = _FlyteClientManager(_platform_config.URL.get(), insecure=_platform_config.INSECURE.get()).client - execution_data = client.get_task_execution_data(self.sdk_task_execution.id) - - # Inputs are returned inline unless they are too big, in which case a url blob pointing to them is returned. - if bool(execution_data.full_outputs.literals): - return execution_data.full_outputs - - if execution_data.outputs.bytes > 0: - with _common_utils.AutoDeletingTempDir() as t: - tmp_name = _os.path.join(t.name, "outputs.pb") - _data_proxy.Data.get_data(execution_data.outputs.url, tmp_name) - return _literals.LiteralMap.from_flyte_idl( - _common_utils.load_proto_from_file(_literals_pb2.LiteralMap, tmp_name) - ) - return _literals.LiteralMap({}) - - @_deprecated( - reason="Objects should access client directly, will be removed by 1.0", - version="0.13.0", - ) - def sync(self): - """ - :rtype: None - """ - client = _FlyteClientManager(_platform_config.URL.get(), insecure=_platform_config.INSECURE.get()).client - self.sdk_task_execution._closure = client.get_task_execution(self.sdk_task_execution.id).closure - - @_deprecated( - reason="Objects should access client directly, will be removed by 1.0", - version="0.13.0", - ) - def get_child_executions(self, filters=None): - """ - :param list[flytekit.models.filters.Filter] filters: - :rtype: dict[Text, flytekit.common.nodes.SdkNodeExecution] - """ - client = _FlyteClientManager(_platform_config.URL.get(), insecure=_platform_config.INSECURE.get()).client - return { - v.id.node_id: v - for v in _iterate_node_executions( - client, - task_execution_identifier=self.sdk_task_execution.id, - filters=filters, - ) - } diff --git a/flytekit/engines/loader.py b/flytekit/engines/loader.py deleted file mode 100644 index 336d5b9de6..0000000000 --- a/flytekit/engines/loader.py +++ /dev/null @@ -1,44 +0,0 @@ -import importlib as _importlib - -from flytekit.common.exceptions import scopes as _exception_scopes -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.configuration import sdk as _sdk_config - -_ENGINE_NAME_TO_MODULES_CACHE = { - "flyte": ("flytekit.engines.flyte.engine", "FlyteEngineFactory", None), - "unit": ("flytekit.engines.unit.engine", "UnitTestEngineFactory", None), - # 'local': ('flytekit.engines.local.engine', 'EngineObjectFactory', None) -} - - -def get_engine(engine_name=None): - """ - :param Text engine_name: - :rtype: flytekit.engines.common.BaseExecutionEngineFactory - """ - engine_name = engine_name or _sdk_config.EXECUTION_ENGINE.get() - - # TODO: Allow users to plug-in their own engine code via a config - if engine_name not in _ENGINE_NAME_TO_MODULES_CACHE: - raise _user_exceptions.FlyteValueException( - engine_name, - "Could not load an engine with the identifier '{}'. Known engines are: {}".format( - engine_name, list(_ENGINE_NAME_TO_MODULES_CACHE.keys()) - ), - ) - - module_path, attr, engine_impl = _ENGINE_NAME_TO_MODULES_CACHE[engine_name] - if engine_impl is None: - module = _exception_scopes.user_entry_point(_importlib.import_module)(module_path) - - if not hasattr(module, attr): - raise _user_exceptions.FlyteValueException( - module, - "Failed to load the engine because the attribute named '{}' could not be found" - "in the module '{}'.".format(attr, module_path), - ) - - engine_impl = getattr(module, attr)() - _ENGINE_NAME_TO_MODULES_CACHE[engine_name] = (module_path, attr, engine_impl) - - return engine_impl diff --git a/flytekit/engines/unit/__init__.py b/flytekit/engines/unit/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/flytekit/engines/unit/engine.py b/flytekit/engines/unit/engine.py deleted file mode 100644 index 2e3c4e187b..0000000000 --- a/flytekit/engines/unit/engine.py +++ /dev/null @@ -1,329 +0,0 @@ -import logging as _logging -import os as _os -from datetime import datetime as _datetime - -import six as _six -from flyteidl.plugins import qubole_pb2 as _qubole_pb2 -from google.protobuf.json_format import ParseDict as _ParseDict -from six import moves as _six_moves - -from flytekit.common import constants as _sdk_constants -from flytekit.common import utils as _common_utils -from flytekit.common.exceptions import system as _system_exception -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.types import helpers as _type_helpers -from flytekit.configuration import TemporaryConfiguration as _TemporaryConfiguration -from flytekit.engines import common as _common_engine -from flytekit.engines.unit.mock_stats import MockStats -from flytekit.interfaces.data import data_proxy as _data_proxy -from flytekit.models import array_job as _array_job -from flytekit.models import literals as _literals -from flytekit.models import qubole as _qubole_models -from flytekit.models.core.identifier import WorkflowExecutionIdentifier - - -class UnitTestEngineFactory(_common_engine.BaseExecutionEngineFactory): - def get_task(self, sdk_task): - """ - :param flytekit.common.tasks.task.SdkTask sdk_task: - :rtype: UnitTestEngineTask - """ - if sdk_task.type in { - _sdk_constants.SdkTaskType.PYTHON_TASK, - _sdk_constants.SdkTaskType.SPARK_TASK, - _sdk_constants.SdkTaskType.SENSOR_TASK, - }: - return ReturnOutputsTask(sdk_task) - elif sdk_task.type in { - _sdk_constants.SdkTaskType.DYNAMIC_TASK, - }: - return DynamicTask(sdk_task) - elif sdk_task.type in { - _sdk_constants.SdkTaskType.BATCH_HIVE_TASK, - }: - return HiveTask(sdk_task) - else: - raise _user_exceptions.FlyteAssertion( - "Unit tests are not currently supported for tasks of type: {}".format(sdk_task.type) - ) - - def get_workflow(self, _): - raise _user_exceptions.FlyteAssertion("Unit testing of workflows is not currently supported") - - def get_launch_plan(self, _): - raise _user_exceptions.FlyteAssertion("Unit testing of launch plans is not currently supported") - - def get_task_execution(self, _): - raise _user_exceptions.FlyteAssertion("Unit testing does not return execution handles.") - - def get_node_execution(self, _): - raise _user_exceptions.FlyteAssertion("Unit testing does not return execution handles.") - - def get_workflow_execution(self, _): - raise _user_exceptions.FlyteAssertion("Unit testing does not return execution handles.") - - def fetch_workflow_execution(self, _): - raise _user_exceptions.FlyteAssertion("Unit testing does not fetch execution handles.") - - def fetch_task(self, _): - raise _user_exceptions.FlyteAssertion("Unit testing does not fetch real tasks.") - - def fetch_latest_task(self, named_task): - raise _user_exceptions.FlyteAssertion("Unit testing does not fetch the real latest task.") - - def fetch_launch_plan(self, _): - raise _user_exceptions.FlyteAssertion("Unit testing does not fetch real launch plans.") - - def fetch_workflow(self, _): - raise _user_exceptions.FlyteAssertion("Unit testing does not fetch real workflows.") - - -class UnitTestEngineTask(_common_engine.BaseTaskExecutor): - def execute(self, inputs, context=None): - """ - Just execute the function and return the outputs as a user-readable dictionary. - :param flytekit.models.literals.LiteralMap inputs: - :param context: - :rtype: dict[Text,flytekit.models.common.FlyteIdlEntity] - """ - with _TemporaryConfiguration( - _os.path.join(_os.path.dirname(__file__), "unit.config"), - internal_overrides={"image": "unit_image"}, - ): - with _common_utils.AutoDeletingTempDir("unit_test_dir") as working_directory: - with _data_proxy.LocalWorkingDirectoryContext(working_directory): - return self._transform_for_user_output(self._execute_user_code(inputs)) - - def _execute_user_code(self, inputs): - """ - :param flytekit.models.literals.LiteralMap inputs: - :rtype: dict[Text,flytekit.models.common.FlyteIdlEntity] - """ - with _common_utils.AutoDeletingTempDir("user_dir") as user_working_directory: - return self.sdk_task.execute( - _common_engine.EngineContext( - execution_id=WorkflowExecutionIdentifier(project="unit_test", domain="unit_test", name="unit_test"), - execution_date=_datetime.utcnow(), - stats=MockStats(), - logging=_logging, # TODO: A mock logging object that we can read later. - tmp_dir=user_working_directory, - ), - inputs, - ) - - def _transform_for_user_output(self, outputs): - """ - Take whatever is returned from the task execution and convert to a reasonable output for the behavior of this - task's unit test. - :param dict[Text,flytekit.models.common.FlyteIdlEntity] outputs: - :rtype: T - """ - return outputs - - def register(self, identifier, version): - raise _user_exceptions.FlyteAssertion("You cannot register unit test tasks.") - - def launch( - self, - project, - domain, - name=None, - inputs=None, - notification_overrides=None, - label_overrides=None, - annotation_overrides=None, - auth_role=None, - ): - raise _user_exceptions.FlyteAssertion("You cannot launch unit test tasks.") - - -class ReturnOutputsTask(UnitTestEngineTask): - def _transform_for_user_output(self, outputs): - """ - Just return the outputs as a user-readable dictionary. - :param dict[Text,flytekit.models.common.FlyteIdlEntity] outputs: - :rtype: T - """ - literal_map = outputs[_sdk_constants.OUTPUT_FILE_NAME] - return { - name: _type_helpers.get_sdk_value_from_literal( - literal_map.literals[name], - sdk_type=_type_helpers.get_sdk_type_from_literal_type(variable.type), - ).to_python_std() - for name, variable in _six.iteritems(self.sdk_task.interface.outputs) - } - - -class DynamicTask(ReturnOutputsTask): - def __init__(self, *args, **kwargs): - self._has_workflow_node = False - super(DynamicTask, self).__init__(*args, **kwargs) - - def _transform_for_user_output(self, outputs): - if self.has_workflow_node: - # If a workflow node has been detected, then we skip any transformation - # This is to support the early termination behavior of the unit test engine when it comes to dynamic tasks - # that produce launch plan or subworkflow nodes. - # See the warning message in the code below for additional information - return outputs - return super(DynamicTask, self)._transform_for_user_output(outputs) - - def _execute_user_code(self, inputs): - """ - :param flytekit.models.literals.LiteralMap inputs: - :rtype: dict[Text,flytekit.models.common.FlyteIdlEntity] - """ - results = super(DynamicTask, self)._execute_user_code(inputs) - if _sdk_constants.FUTURES_FILE_NAME in results: - futures = results[_sdk_constants.FUTURES_FILE_NAME] - sub_task_outputs = {} - tasks_map = {task.id: task for task in futures.tasks} - - for future_node in futures.nodes: - if future_node.workflow_node is not None: - # TODO: implement proper unit testing for launchplan and subworkflow nodes somehow - _logging.warning( - "A workflow node has been detected in the output of the dynamic task. The " - "Flytekit unit test engine is incomplete for dynamic tasks that return launch " - "plans or subworkflows. The generated dynamic job spec will be returned but " - "they will not be run." - ) - # For now, just return the output of the parent task - self._has_workflow_node = True - return results - task = tasks_map[future_node.task_node.reference_id] - if task.type == _sdk_constants.SdkTaskType.CONTAINER_ARRAY_TASK: - sub_task_output = DynamicTask.execute_array_task(future_node.id, task, results) - elif task.type == _sdk_constants.SdkTaskType.SPARK_TASK: - # This is required because `_transform_for_user_output` function about is invoked which - # checks for outputs - self._has_workflow_node = True - return results - elif task.type == _sdk_constants.SdkTaskType.HIVE_JOB: - # TODO: futures.outputs should have the Schema instances. - # After schema is implemented, fill out random data into the random locations - # then check output in test function - # Even though we recommend people use typed schemas, they might not always do so... - # in which case it'll be impossible to predict the actual schema, we should support a - # way for unit test authors to provide fake data regardless - sub_task_output = None - else: - inputs_path = _os.path.join(future_node.id, _sdk_constants.INPUT_FILE_NAME) - if inputs_path not in results: - raise _system_exception.FlyteSystemAssertion( - "dynamic task hasn't generated expected inputs document [{}] found {}".format( - future_node.id, list(results.keys()) - ) - ) - sub_task_output = UnitTestEngineFactory().get_task(task).execute(results[inputs_path]) - sub_task_outputs[future_node.id] = sub_task_output - - results[_sdk_constants.OUTPUT_FILE_NAME] = _literals.LiteralMap( - literals={ - binding.var: DynamicTask.fulfil_bindings(binding.binding, sub_task_outputs) - for binding in futures.outputs - } - ) - return results - - @property - def has_workflow_node(self): - """ - :rtype: bool - """ - return self._has_workflow_node - - @staticmethod - def execute_array_task(root_input_path, task, array_inputs): - array_job = _array_job.ArrayJob.from_dict(task.custom) - outputs = {} - for job_index in _six_moves.range(0, array_job.size): - inputs_path = _os.path.join( - root_input_path, - _six.text_type(job_index), - _sdk_constants.INPUT_FILE_NAME, - ) - if inputs_path not in array_inputs: - raise _system_exception.FlyteSystemAssertion( - "dynamic task hasn't generated expected inputs document [{}].".format(inputs_path) - ) - - input_proto = array_inputs[inputs_path] - # All outputs generated by the same array job will have the same key in sub_task_outputs, - # they will, however, differ in the var names; they will be on the format []. - # e.g. [1].out1 - for key, val in _six.iteritems( - ReturnOutputsTask( - task.assign_type_and_return(_sdk_constants.SdkTaskType.PYTHON_TASK) # TODO: This is weird - ).execute(input_proto) - ): - outputs["[{}].{}".format(job_index, key)] = val - return outputs - - @staticmethod - def fulfil_bindings(binding_data, fulfilled_promises): - """ - Substitutes promise values in binding_data with model Literal values built from python std values in - fulfilled_promises - - :param _interface.BindingData binding_data: - :param dict[Text,T] fulfilled_promises: - :rtype: - """ - if binding_data.scalar: - return _literals.Literal(scalar=binding_data.scalar) - elif binding_data.collection: - return _literals.Literal( - collection=_literals.LiteralCollection( - [ - DynamicTask.fulfil_bindings(sub_binding_data, fulfilled_promises) - for sub_binding_data in binding_data.collection.bindings - ] - ) - ) - elif binding_data.promise: - if binding_data.promise.node_id not in fulfilled_promises: - raise _system_exception.FlyteSystemAssertion( - "Expecting output of node [{}] but that hasn't been produced.".format(binding_data.promise.node_id) - ) - node_output = fulfilled_promises[binding_data.promise.node_id] - if binding_data.promise.var not in node_output: - raise _system_exception.FlyteSystemAssertion( - "Expecting output [{}] of node [{}] but that hasn't been produced.".format( - binding_data.promise.var, binding_data.promise.node_id - ) - ) - - return binding_data.promise.sdk_type.from_python_std(node_output[binding_data.promise.var]) - elif binding_data.map: - return _literals.Literal( - map=_literals.LiteralMap( - { - k: DynamicTask.fulfil_bindings(sub_binding_data, fulfilled_promises) - for k, sub_binding_data in _six.iteritems(binding_data.map.bindings) - } - ) - ) - - -class HiveTask(DynamicTask): - def _transform_for_user_output(self, outputs): - """ - Just execute the function and return the list of Hive queries returned. - :param dict[Text,flytekit.models.common.FlyteIdlEntity] outputs: - :rtype: list[Text] - """ - futures = outputs.get(_sdk_constants.FUTURES_FILE_NAME) - if futures: - queries = [] - task_ids_to_defs = { - t.id.name: _qubole_models.QuboleHiveJob.from_flyte_idl( - _ParseDict(t.custom, _qubole_pb2.QuboleHiveJob()) - ) - for t in futures.tasks - } - for node in futures.nodes: - queries.append(task_ids_to_defs[node.task_node.reference_id.name].query.query) - return queries - else: - return [] diff --git a/flytekit/engines/unit/unit.config b/flytekit/engines/unit/unit.config deleted file mode 100644 index 8e2c00a83e..0000000000 --- a/flytekit/engines/unit/unit.config +++ /dev/null @@ -1,14 +0,0 @@ -[sdk] - -workflow_packages=this.module,that.module - -[auth] -assumable_iam_role=unit_test_role - -[container] - -image=some_docker_repo:some_image_name:tag - -[platform] - -url=unittest diff --git a/flytekit/common/__init__.py b/flytekit/exceptions/__init__.py similarity index 100% rename from flytekit/common/__init__.py rename to flytekit/exceptions/__init__.py diff --git a/flytekit/common/exceptions/base.py b/flytekit/exceptions/base.py similarity index 100% rename from flytekit/common/exceptions/base.py rename to flytekit/exceptions/base.py diff --git a/flytekit/common/exceptions/scopes.py b/flytekit/exceptions/scopes.py similarity index 97% rename from flytekit/common/exceptions/scopes.py rename to flytekit/exceptions/scopes.py index a32d9dbcc6..60a4afa97e 100644 --- a/flytekit/common/exceptions/scopes.py +++ b/flytekit/exceptions/scopes.py @@ -3,9 +3,9 @@ from wrapt import decorator as _decorator -from flytekit.common.exceptions import base as _base_exceptions -from flytekit.common.exceptions import system as _system_exceptions -from flytekit.common.exceptions import user as _user_exceptions +from flytekit.exceptions import base as _base_exceptions +from flytekit.exceptions import system as _system_exceptions +from flytekit.exceptions import user as _user_exceptions from flytekit.models.core import errors as _error_model diff --git a/flytekit/common/exceptions/system.py b/flytekit/exceptions/system.py similarity index 95% rename from flytekit/common/exceptions/system.py rename to flytekit/exceptions/system.py index 5802279590..63c43e8879 100644 --- a/flytekit/common/exceptions/system.py +++ b/flytekit/exceptions/system.py @@ -1,4 +1,4 @@ -from flytekit.common.exceptions import base as _base_exceptions +from flytekit.exceptions import base as _base_exceptions class FlyteSystemException(_base_exceptions.FlyteRecoverableException): diff --git a/flytekit/common/exceptions/user.py b/flytekit/exceptions/user.py similarity index 94% rename from flytekit/common/exceptions/user.py rename to flytekit/exceptions/user.py index acb5dd7997..a93510cae1 100644 --- a/flytekit/common/exceptions/user.py +++ b/flytekit/exceptions/user.py @@ -1,5 +1,5 @@ -from flytekit.common.exceptions.base import FlyteException as _FlyteException -from flytekit.common.exceptions.base import FlyteRecoverableException as _Recoverable +from flytekit.exceptions.base import FlyteException as _FlyteException +from flytekit.exceptions.base import FlyteRecoverableException as _Recoverable class FlyteUserException(_FlyteException): diff --git a/flytekit/extend/__init__.py b/flytekit/extend/__init__.py index 2c22fd5bd9..d420310fa2 100644 --- a/flytekit/extend/__init__.py +++ b/flytekit/extend/__init__.py @@ -33,12 +33,11 @@ DataPersistencePlugins """ -from flytekit.common.translator import get_serializable from flytekit.core import context_manager from flytekit.core.base_sql_task import SQLTask from flytekit.core.base_task import IgnoreOutputs, PythonTask, TaskResolverMixin from flytekit.core.class_based_resolver import ClassStorageTaskResolver -from flytekit.core.context_manager import ExecutionState, Image, ImageConfig, SerializationSettings +from flytekit.core.context_manager import ExecutionState, Image, ImageConfig, SecretsManager, SerializationSettings from flytekit.core.data_persistence import DataPersistence, DataPersistencePlugins from flytekit.core.interface import Interface from flytekit.core.promise import Promise @@ -46,3 +45,4 @@ from flytekit.core.shim_task import ExecutableTemplateShimTask, ShimTaskExecutor from flytekit.core.task import TaskPlugins from flytekit.core.type_engine import DictTransformer, T, TypeEngine, TypeTransformer +from flytekit.tools.translator import get_serializable diff --git a/flytekit/extras/persistence/gcs_gsutil.py b/flytekit/extras/persistence/gcs_gsutil.py index bb4ec31487..7e7711d64a 100644 --- a/flytekit/extras/persistence/gcs_gsutil.py +++ b/flytekit/extras/persistence/gcs_gsutil.py @@ -2,9 +2,9 @@ import typing from shutil import which as shell_which -from flytekit.common.exceptions.user import FlyteUserException from flytekit.configuration import gcp from flytekit.core.data_persistence import DataPersistence, DataPersistencePlugins +from flytekit.exceptions.user import FlyteUserException from flytekit.tools import subprocess diff --git a/flytekit/extras/persistence/http.py b/flytekit/extras/persistence/http.py index d9fa4674d7..30fa8d0f65 100644 --- a/flytekit/extras/persistence/http.py +++ b/flytekit/extras/persistence/http.py @@ -3,8 +3,8 @@ import requests -from flytekit.common.exceptions import user from flytekit.core.data_persistence import DataPersistence, DataPersistencePlugins +from flytekit.exceptions import user from flytekit.loggers import logger diff --git a/flytekit/extras/persistence/s3_awscli.py b/flytekit/extras/persistence/s3_awscli.py index ddf26ec4d5..3b24fef94b 100644 --- a/flytekit/extras/persistence/s3_awscli.py +++ b/flytekit/extras/persistence/s3_awscli.py @@ -7,9 +7,9 @@ from shutil import which as shell_which from typing import Dict, List, Optional -from flytekit.common.exceptions.user import FlyteUserException from flytekit.configuration import aws from flytekit.core.data_persistence import DataPersistence, DataPersistencePlugins +from flytekit.exceptions.user import FlyteUserException from flytekit.tools import subprocess diff --git a/flytekit/common/core/identifier.py b/flytekit/interfaces/cli_identifiers.py similarity index 89% rename from flytekit/common/core/identifier.py rename to flytekit/interfaces/cli_identifiers.py index c7b12a5190..18cfad424b 100644 --- a/flytekit/common/core/identifier.py +++ b/flytekit/interfaces/cli_identifiers.py @@ -1,18 +1,15 @@ -import six as _six - -from flytekit.common import sdk_bases as _sdk_bases -from flytekit.common.exceptions import user as _user_exceptions +from flytekit.exceptions import user as _user_exceptions from flytekit.models.core import identifier as _core_identifier -class Identifier(_core_identifier.Identifier, metaclass=_sdk_bases.ExtendedSdkType): +class Identifier(_core_identifier.Identifier): _STRING_TO_TYPE_MAP = { "lp": _core_identifier.ResourceType.LAUNCH_PLAN, "wf": _core_identifier.ResourceType.WORKFLOW, "tsk": _core_identifier.ResourceType.TASK, } - _TYPE_TO_STRING_MAP = {v: k for k, v in _six.iteritems(_STRING_TO_TYPE_MAP)} + _TYPE_TO_STRING_MAP = {v: k for k, v in _STRING_TO_TYPE_MAP.items()} @classmethod def promote_from_model(cls, base_model): @@ -28,6 +25,11 @@ def promote_from_model(cls, base_model): base_model.version, ) + @classmethod + def from_flyte_idl(cls, pb2_object): + base_model = super().from_flyte_idl(pb2_object) + return cls.promote_from_model(base_model) + @classmethod def from_python_std(cls, string): """ @@ -63,106 +65,116 @@ def __str__(self): ) -class WorkflowExecutionIdentifier(_core_identifier.WorkflowExecutionIdentifier, metaclass=_sdk_bases.ExtendedSdkType): +class TaskExecutionIdentifier(_core_identifier.TaskExecutionIdentifier): @classmethod def promote_from_model(cls, base_model): """ - :param flytekit.models.core.identifier.WorkflowExecutionIdentifier base_model: - :rtype: WorkflowExecutionIdentifier + :param flytekit.models.core.identifier.TaskExecutionIdentifier base_model: + :rtype: TaskExecutionIdentifier """ return cls( - base_model.project, - base_model.domain, - base_model.name, + task_id=base_model.task_id, + node_execution_id=base_model.node_execution_id, + retry_attempt=base_model.retry_attempt, ) + @classmethod + def from_flyte_idl(cls, pb2_object): + base_model = super().from_flyte_idl(pb2_object) + return cls.promote_from_model(base_model) + @classmethod def from_python_std(cls, string): """ Parses a string in the correct format into an identifier :param Text string: - :rtype: WorkflowExecutionIdentifier + :rtype: TaskExecutionIdentifier """ segments = string.split(":") - if len(segments) != 4: + if len(segments) != 10: raise _user_exceptions.FlyteValueException( string, "The provided string was not in a parseable format. The string for an identifier must be in the format" - " ex:project:domain:name.", + " te:exec_project:exec_domain:exec_name:node_id:task_project:task_domain:task_name:task_version:retry.", ) - resource_type, project, domain, name = segments + resource_type, ep, ed, en, node_id, tp, td, tn, tv, retry = segments - if resource_type != "ex": + if resource_type != "te": raise _user_exceptions.FlyteValueException( resource_type, "The provided string could not be parsed. The first element of an execution identifier must be 'ex'.", ) return cls( - project, - domain, - name, + task_id=Identifier(_core_identifier.ResourceType.TASK, tp, td, tn, tv), + node_execution_id=_core_identifier.NodeExecutionIdentifier( + node_id=node_id, + execution_id=_core_identifier.WorkflowExecutionIdentifier(ep, ed, en), + ), + retry_attempt=int(retry), ) def __str__(self): - return "ex:{}:{}:{}".format(self.project, self.domain, self.name) + return "te:{ep}:{ed}:{en}:{node_id}:{tp}:{td}:{tn}:{tv}:{retry}".format( + ep=self.node_execution_id.execution_id.project, + ed=self.node_execution_id.execution_id.domain, + en=self.node_execution_id.execution_id.name, + node_id=self.node_execution_id.node_id, + tp=self.task_id.project, + td=self.task_id.domain, + tn=self.task_id.name, + tv=self.task_id.version, + retry=self.retry_attempt, + ) -class TaskExecutionIdentifier(_core_identifier.TaskExecutionIdentifier, metaclass=_sdk_bases.ExtendedSdkType): +class WorkflowExecutionIdentifier(_core_identifier.WorkflowExecutionIdentifier): @classmethod def promote_from_model(cls, base_model): """ - :param flytekit.models.core.identifier.TaskExecutionIdentifier base_model: - :rtype: TaskExecutionIdentifier + :param flytekit.models.core.identifier.WorkflowExecutionIdentifier base_model: + :rtype: WorkflowExecutionIdentifier """ return cls( - task_id=base_model.task_id, - node_execution_id=base_model.node_execution_id, - retry_attempt=base_model.retry_attempt, + base_model.project, + base_model.domain, + base_model.name, ) + @classmethod + def from_flyte_idl(cls, pb2_object): + base_model = super().from_flyte_idl(pb2_object) + return cls.promote_from_model(base_model) + @classmethod def from_python_std(cls, string): """ Parses a string in the correct format into an identifier :param Text string: - :rtype: TaskExecutionIdentifier + :rtype: WorkflowExecutionIdentifier """ segments = string.split(":") - if len(segments) != 10: + if len(segments) != 4: raise _user_exceptions.FlyteValueException( string, "The provided string was not in a parseable format. The string for an identifier must be in the format" - " te:exec_project:exec_domain:exec_name:node_id:task_project:task_domain:task_name:task_version:retry.", + " ex:project:domain:name.", ) - resource_type, ep, ed, en, node_id, tp, td, tn, tv, retry = segments + resource_type, project, domain, name = segments - if resource_type != "te": + if resource_type != "ex": raise _user_exceptions.FlyteValueException( resource_type, "The provided string could not be parsed. The first element of an execution identifier must be 'ex'.", ) return cls( - task_id=Identifier(_core_identifier.ResourceType.TASK, tp, td, tn, tv), - node_execution_id=_core_identifier.NodeExecutionIdentifier( - node_id=node_id, - execution_id=_core_identifier.WorkflowExecutionIdentifier(ep, ed, en), - ), - retry_attempt=int(retry), + project, + domain, + name, ) def __str__(self): - return "te:{ep}:{ed}:{en}:{node_id}:{tp}:{td}:{tn}:{tv}:{retry}".format( - ep=self.node_execution_id.execution_id.project, - ed=self.node_execution_id.execution_id.domain, - en=self.node_execution_id.execution_id.name, - node_id=self.node_execution_id.node_id, - tp=self.task_id.project, - td=self.task_id.domain, - tn=self.task_id.name, - tv=self.task_id.version, - retry=self.retry_attempt, - ) + return "ex:{}:{}:{}".format(self.project, self.domain, self.name) diff --git a/flytekit/interfaces/data/__init__.py b/flytekit/interfaces/data/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/flytekit/interfaces/data/common.py b/flytekit/interfaces/data/common.py deleted file mode 100644 index 1544b608b6..0000000000 --- a/flytekit/interfaces/data/common.py +++ /dev/null @@ -1,54 +0,0 @@ -class DataProxy(object): - def __init__(self, name: str): - self._name = name - - @property - def name(self) -> str: - return self._name - - def exists(self, path): - """ - :param path: - :rtype: bool: whether the file exists or not - """ - pass - - def download_directory(self, remote_path, local_path): - """ - :param Text remote_path: - :param Text local_path: - """ - pass - - def download(self, remote_path, local_path): - """ - :param Text remote_path: - :param Text local_path: - """ - pass - - def upload(self, file_path, to_path): - """ - :param Text file_path: - :param Text to_path: - """ - pass - - def upload_directory(self, local_path, remote_path): - """ - :param Text local_path: - :param Text remote_path: - """ - pass - - def get_random_path(self): - """ - :rtype: Text - """ - pass - - def get_random_directory(self): - """ - :rtype: Text - """ - pass diff --git a/flytekit/interfaces/data/data_proxy.py b/flytekit/interfaces/data/data_proxy.py deleted file mode 100644 index 7cf5dcae58..0000000000 --- a/flytekit/interfaces/data/data_proxy.py +++ /dev/null @@ -1,172 +0,0 @@ -from flytekit.common import constants as _constants -from flytekit.common import utils as _common_utils -from flytekit.common.exceptions import user as _user_exception -from flytekit.configuration import platform as _platform_config -from flytekit.configuration import sdk as _sdk_config -from flytekit.interfaces.data.gcs import gcs_proxy as _gcs_proxy -from flytekit.interfaces.data.http import http_data_proxy as _http_data_proxy -from flytekit.interfaces.data.local import local_file_proxy as _local_file_proxy -from flytekit.interfaces.data.s3 import s3proxy as _s3proxy - - -class LocalWorkingDirectoryContext(object): - _CONTEXTS = [] - - def __init__(self, directory): - self._directory = directory - - def __enter__(self): - self._CONTEXTS.append(self._directory) - - def __exit__(self, exc_type, exc_val, exc_tb): - self._CONTEXTS.pop() - - @classmethod - def get(cls): - return cls._CONTEXTS[-1] if cls._CONTEXTS else None - - -class _OutputDataContext(object): - _CONTEXTS = [_local_file_proxy.LocalFileProxy(_sdk_config.LOCAL_SANDBOX.get())] - - def __init__(self, context): - self._context = context - - def __enter__(self): - self._CONTEXTS.append(self._context) - - def __exit__(self, exc_type, exc_val, exc_tb): - self._CONTEXTS.pop() - - @classmethod - def get_active_proxy(cls): - return cls._CONTEXTS[-1] - - @classmethod - def get_default_proxy(cls): - return cls._CONTEXTS[0] - - -class LocalDataContext(_OutputDataContext): - def __init__(self, sandbox): - """ - :param Text sandbox: - """ - super(LocalDataContext, self).__init__(_local_file_proxy.LocalFileProxy(sandbox)) - - -class RemoteDataContext(_OutputDataContext): - _CLOUD_PROVIDER_TO_PROXIES = { - _constants.CloudProvider.AWS: _s3proxy.AwsS3Proxy, - _constants.CloudProvider.GCP: _gcs_proxy.GCSProxy, - } - - def __init__(self, cloud_provider=None, raw_output_data_prefix_override=None): - """ - :param Optional[Text] cloud_provider: From flytekit.common.constants.CloudProvider enum - """ - cloud_provider = cloud_provider or _platform_config.CLOUD_PROVIDER.get() - proxy_class = type(self)._CLOUD_PROVIDER_TO_PROXIES.get(cloud_provider, None) - if proxy_class is None: - raise _user_exception.FlyteAssertion( - "Configured cloud provider is not supported for data I/O. Received: {}, expected one of: {}".format( - cloud_provider, list(type(self)._CLOUD_PROVIDER_TO_PROXIES.keys()) - ) - ) - proxy = proxy_class(raw_output_data_prefix_override) - super(RemoteDataContext, self).__init__(proxy) - - -class Data(object): - # TODO: More proxies for more environments. - _DATA_PROXIES = { - "s3:/": _s3proxy.AwsS3Proxy(), - "gs:/": _gcs_proxy.GCSProxy(), - "http://": _http_data_proxy.HttpFileProxy(), - "https://": _http_data_proxy.HttpFileProxy(), - } - - @classmethod - def _load_data_proxy_by_path(cls, path): - """ - :param Text path: - :rtype: flytekit.interfaces.data.common.DataProxy - """ - for k, v in cls._DATA_PROXIES.items(): - if path.startswith(k): - return v - return _OutputDataContext.get_default_proxy() - - @classmethod - def data_exists(cls, path): - """ - :param Text path: - :rtype: bool: whether the file exists or not - """ - with _common_utils.PerformanceTimer("Check file exists {}".format(path)): - proxy = cls._load_data_proxy_by_path(path) - return proxy.exists(path) - - @classmethod - def get_data(cls, remote_path, local_path, is_multipart=False): - """ - :param Text remote_path: - :param Text local_path: - :param bool is_multipart: - """ - try: - with _common_utils.PerformanceTimer("Copying ({} -> {})".format(remote_path, local_path)): - proxy = cls._load_data_proxy_by_path(remote_path) - if is_multipart: - proxy.download_directory(remote_path, local_path) - else: - proxy.download(remote_path, local_path) - except Exception as ex: - raise _user_exception.FlyteAssertion( - "Failed to get data from {remote_path} to {local_path} (recursive={is_multipart}).\n\n" - "Original exception: {error_string}".format( - remote_path=remote_path, - local_path=local_path, - is_multipart=is_multipart, - error_string=str(ex), - ) - ) - - @classmethod - def put_data(cls, local_path, remote_path, is_multipart=False): - """ - :param Text local_path: - :param Text remote_path: - :param bool is_multipart: - """ - try: - with _common_utils.PerformanceTimer("Writing ({} -> {})".format(local_path, remote_path)): - proxy = cls._load_data_proxy_by_path(remote_path) - if is_multipart: - proxy.upload_directory(local_path, remote_path) - else: - proxy.upload(local_path, remote_path) - except Exception as ex: - raise _user_exception.FlyteAssertion( - "Failed to put data from {local_path} to {remote_path} (recursive={is_multipart}).\n\n" - "Original exception: {error_string}".format( - remote_path=remote_path, - local_path=local_path, - is_multipart=is_multipart, - error_string=str(ex), - ) - ) - - @classmethod - def get_remote_path(cls): - """ - :rtype: Text - """ - return _OutputDataContext.get_active_proxy().get_random_path() - - @classmethod - def get_remote_directory(cls): - """ - :rtype: Text - """ - return _OutputDataContext.get_active_proxy().get_random_directory() diff --git a/flytekit/interfaces/data/gcs/__init__.py b/flytekit/interfaces/data/gcs/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/flytekit/interfaces/data/gcs/gcs_proxy.py b/flytekit/interfaces/data/gcs/gcs_proxy.py deleted file mode 100644 index e8440299cd..0000000000 --- a/flytekit/interfaces/data/gcs/gcs_proxy.py +++ /dev/null @@ -1,150 +0,0 @@ -import os as _os -import sys as _sys -import uuid as _uuid - -from flytekit.common.exceptions.user import FlyteUserException as _FlyteUserException -from flytekit.configuration import gcp as _gcp_config -from flytekit.interfaces import random as _flyte_random -from flytekit.interfaces.data import common as _common_data -from flytekit.tools import subprocess as _subprocess - -if _sys.version_info >= (3,): - from shutil import which as _which -else: - from distutils.spawn import find_executable as _which - - -def _update_cmd_config_and_execute(cmd): - env = _os.environ.copy() - return _subprocess.check_call(cmd, env=env) - - -def _amend_path(path): - return _os.path.join(path, "*") if not path.endswith("*") else path - - -class GCSProxy(_common_data.DataProxy): - _GS_UTIL_CLI = "gsutil" - - def __init__(self, raw_output_data_prefix_override: str = None): - """ - :param raw_output_data_prefix_override: Instead of relying on the AWS or GCS configuration (see - S3_SHARD_FORMATTER for AWS and GCS_PREFIX for GCP) setting when computing the shard - path (_get_shard_path), use this prefix instead as a base. This code assumes that the - path passed in is correct. That is, an S3 path won't be passed in when running on GCP. - """ - self._raw_output_data_prefix_override = raw_output_data_prefix_override - super(GCSProxy, self).__init__(name="gcs-gsutil") - - @property - def raw_output_data_prefix_override(self) -> str: - return self._raw_output_data_prefix_override - - @staticmethod - def _check_binary(): - """ - Make sure that the `gsutil` cli is present - """ - if not _which(GCSProxy._GS_UTIL_CLI): - raise _FlyteUserException("gsutil (gcloud cli) not found! Please install.") - - @staticmethod - def _maybe_with_gsutil_parallelism(*gsutil_args): - """ - Check if we should run `gsutil` with the `-m` flag that enables - parallelism via multiple threads/processes. Additional tweaking of - this behavior can be achieved via the .boto configuration file. See: - https://cloud.google.com/storage/docs/boto-gsutil - """ - cmd = [GCSProxy._GS_UTIL_CLI] - if _gcp_config.GSUTIL_PARALLELISM.get(): - cmd.append("-m") - cmd.extend(gsutil_args) - - return cmd - - def exists(self, remote_path): - """ - :param Text remote_path: remote gs:// path - :rtype bool: whether the gs file exists or not - """ - GCSProxy._check_binary() - - if not remote_path.startswith("gs://"): - raise ValueError("Not an GS Key. Please use FQN (GS ARN) of the format gs://...") - - cmd = [GCSProxy._GS_UTIL_CLI, "-q", "stat", remote_path] - try: - _update_cmd_config_and_execute(cmd) - return True - except Exception: - return False - - def download_directory(self, remote_path, local_path): - """ - :param Text remote_path: remote gs:// path - :param Text local_path: directory to copy to - """ - GCSProxy._check_binary() - - if not remote_path.startswith("gs://"): - raise ValueError("Not an GS Key. Please use FQN (GS ARN) of the format gs://...") - - cmd = self._maybe_with_gsutil_parallelism("cp", "-r", _amend_path(remote_path), local_path) - return _update_cmd_config_and_execute(cmd) - - def download(self, remote_path, local_path): - """ - :param Text remote_path: remote gs:// path - :param Text local_path: directory to copy to - """ - if not remote_path.startswith("gs://"): - raise ValueError("Not an GS Key. Please use FQN (GS ARN) of the format gs://...") - - GCSProxy._check_binary() - - cmd = self._maybe_with_gsutil_parallelism("cp", remote_path, local_path) - return _update_cmd_config_and_execute(cmd) - - def upload(self, file_path, to_path): - """ - :param Text file_path: - :param Text to_path: - """ - GCSProxy._check_binary() - - cmd = self._maybe_with_gsutil_parallelism("cp", file_path, to_path) - return _update_cmd_config_and_execute(cmd) - - def upload_directory(self, local_path, remote_path): - """ - :param Text local_path: - :param Text remote_path: - """ - if not remote_path.startswith("gs://"): - raise ValueError("Not an GS Key. Please use FQN (GS ARN) of the format gs://...") - - GCSProxy._check_binary() - - cmd = self._maybe_with_gsutil_parallelism( - "cp", - "-r", - _amend_path(local_path), - remote_path if remote_path.endswith("/") else remote_path + "/", - ) - return _update_cmd_config_and_execute(cmd) - - def get_random_path(self) -> str: - """ - If this object was created with a raw output data prefix, usually set by Propeller/Plugins at execution time - and piped all the way here, it will be used instead of referencing the GCS_PREFIX configuration. - """ - key = _uuid.UUID(int=_flyte_random.random.getrandbits(128)).hex - prefix = self.raw_output_data_prefix_override or _gcp_config.GCS_PREFIX.get() - return _os.path.join(prefix, key) - - def get_random_directory(self): - """ - :rtype: Text - """ - return self.get_random_path() + "/" diff --git a/flytekit/interfaces/data/http/__init__.py b/flytekit/interfaces/data/http/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/flytekit/interfaces/data/http/http_data_proxy.py b/flytekit/interfaces/data/http/http_data_proxy.py deleted file mode 100644 index 33e1909906..0000000000 --- a/flytekit/interfaces/data/http/http_data_proxy.py +++ /dev/null @@ -1,80 +0,0 @@ -import requests as _requests - -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.interfaces.data import common as _common_data - - -class HttpFileProxy(_common_data.DataProxy): - - _HTTP_OK = 200 - _HTTP_FORBIDDEN = 403 - _HTTP_NOT_FOUND = 404 - - def __init__(self): - super(HttpFileProxy, self).__init__(name="http") - - def exists(self, path): - """ - :param Text path: the path of the file - :rtype bool: whether the file exists or not - """ - rsp = _requests.head(path) - allowed_codes = { - type(self)._HTTP_OK, - type(self)._HTTP_NOT_FOUND, - type(self)._HTTP_FORBIDDEN, - } - if rsp.status_code not in allowed_codes: - raise _user_exceptions.FlyteValueException( - rsp.status_code, - "Data at {} could not be checked for existence. Expected one of: {}".format(path, allowed_codes), - ) - return rsp.status_code == type(self)._HTTP_OK - - def download_directory(self, from_path, to_path): - """ - :param Text from_path: - :param Text to_path: - """ - raise _user_exceptions.FlyteAssertion("Reading data recursively from HTTP endpoint is not currently supported.") - - def download(self, from_path, to_path): - """ - :param Text from_path: - :param Text to_path: - """ - - rsp = _requests.get(from_path) - if rsp.status_code != type(self)._HTTP_OK: - raise _user_exceptions.FlyteValueException( - rsp.status_code, - "Request for data @ {} failed. Expected status code {}".format(from_path, type(self)._HTTP_OK), - ) - with open(to_path, "wb") as writer: - writer.write(rsp.content) - - def upload(self, from_path, to_path): - """ - :param Text from_path: - :param Text to_path: - """ - raise _user_exceptions.FlyteAssertion("Writing data to HTTP endpoint is not currently supported.") - - def upload_directory(self, from_path, to_path): - """ - :param Text from_path: - :param Text to_path: - """ - raise _user_exceptions.FlyteAssertion("Writing data to HTTP endpoint is not currently supported.") - - def get_random_path(self): - """ - :rtype: Text - """ - raise _user_exceptions.FlyteAssertion("Writing data to HTTP endpoint is not currently supported.") - - def get_random_directory(self): - """ - :rtype: Text - """ - raise _user_exceptions.FlyteAssertion("Writing data to HTTP endpoint is not currently supported.") diff --git a/flytekit/interfaces/data/local/__init__.py b/flytekit/interfaces/data/local/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/flytekit/interfaces/data/local/local_file_proxy.py b/flytekit/interfaces/data/local/local_file_proxy.py deleted file mode 100644 index 2c9ea33cf5..0000000000 --- a/flytekit/interfaces/data/local/local_file_proxy.py +++ /dev/null @@ -1,92 +0,0 @@ -import os as _os -import uuid as _uuid -from distutils import dir_util as _dir_util -from shutil import copyfile as _copyfile - -from flytekit.interfaces import random as _flyte_random -from flytekit.interfaces.data import common as _common_data - - -def _make_local_path(path): - if not _os.path.exists(path): - try: - _os.makedirs(path) - except OSError: # Guard against race condition - if not _os.path.isdir(path): - raise - - -def strip_file_header(path: str) -> str: - if path.startswith("file://"): - return path.replace("file://", "", 1) - return path - - -class LocalFileProxy(_common_data.DataProxy): - def __init__(self, sandbox): - """ - :param Text sandbox: - """ - super().__init__(name="local") - self._sandbox = sandbox - - @property - def sandbox(self) -> str: - return self._sandbox - - def exists(self, path): - """ - :param Text path: the path of the file - :rtype bool: whether the file exists or not - """ - return _os.path.exists(strip_file_header(path)) - - def download_directory(self, from_path, to_path): - """ - :param Text from_path: - :param Text to_path: - """ - if from_path != to_path: - _dir_util.copy_tree(strip_file_header(from_path), strip_file_header(to_path)) - - def download(self, from_path, to_path): - """ - :param Text from_path: - :param Text to_path: - """ - _copyfile(strip_file_header(from_path), strip_file_header(to_path)) - - def upload(self, from_path, to_path): - """ - :param Text from_path: - :param Text to_path: - """ - # Emulate s3's flat storage by automatically creating directory path - _make_local_path(_os.path.dirname(strip_file_header(to_path))) - # Write the object to a local file in the sandbox - _copyfile(strip_file_header(from_path), strip_file_header(to_path)) - - def upload_directory(self, from_path, to_path): - """ - :param Text from_path: - :param Text to_path: - """ - self.download_directory(from_path, to_path) - - def get_random_path(self): - """ - :rtype: Text - """ - # Create a 128-bit random hash because the birthday attack principle shows that there is about a 50% chance of a - # collision between objects when 2^(n/2) objects are created (where n is the number of bits in the hash). - # Assuming Flyte eventually creates 1 trillion pieces of data (~2 ^ 40), the likelihood - # of a collision is 10^-15 with 128-bit...or basically 0. - return _os.path.join(self._sandbox, _uuid.UUID(int=_flyte_random.random.getrandbits(128)).hex) - - def get_random_directory(self): - """ - :rtype: Text - """ - random_dir = self.get_random_path() + "/" - _make_local_path(random_dir) - return random_dir diff --git a/flytekit/interfaces/data/s3/__init__.py b/flytekit/interfaces/data/s3/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/flytekit/interfaces/data/s3/s3proxy.py b/flytekit/interfaces/data/s3/s3proxy.py deleted file mode 100644 index ab11e7b738..0000000000 --- a/flytekit/interfaces/data/s3/s3proxy.py +++ /dev/null @@ -1,224 +0,0 @@ -import logging -import os as _os -import re as _re -import string as _string -import sys as _sys -import time -import uuid as _uuid -from typing import Dict, List - -from six import moves as _six_moves -from six import text_type as _text_type - -from flytekit.common.exceptions.user import FlyteUserException as _FlyteUserException -from flytekit.configuration import aws as _aws_config -from flytekit.interfaces import random as _flyte_random -from flytekit.interfaces.data import common as _common_data -from flytekit.tools import subprocess as _subprocess - -if _sys.version_info >= (3,): - from shutil import which as _which -else: - from distutils.spawn import find_executable as _which - - -def _update_cmd_config_and_execute(cmd: List[str]): - env = _os.environ.copy() - - if _aws_config.ENABLE_DEBUG.get(): - cmd.insert(1, "--debug") - - if _aws_config.S3_ENDPOINT.get() is not None: - cmd.insert(1, _aws_config.S3_ENDPOINT.get()) - cmd.insert(1, _aws_config.S3_ENDPOINT_ARG_NAME) - - if _aws_config.S3_ACCESS_KEY_ID.get() is not None: - env[_aws_config.S3_ACCESS_KEY_ID_ENV_NAME] = _aws_config.S3_ACCESS_KEY_ID.get() - - if _aws_config.S3_SECRET_ACCESS_KEY.get() is not None: - env[_aws_config.S3_SECRET_ACCESS_KEY_ENV_NAME] = _aws_config.S3_SECRET_ACCESS_KEY.get() - - retry = 0 - while True: - try: - return _subprocess.check_call(cmd, env=env) - except Exception as e: - logging.error(f"Exception when trying to execute {cmd}, reason: {str(e)}") - retry += 1 - if retry > _aws_config.RETRIES.get(): - raise - secs = _aws_config.BACKOFF_SECONDS.get() - logging.info(f"Sleeping before retrying again, after {secs} seconds") - time.sleep(secs) - logging.info("Retrying again") - - -def _extra_args(extra_args: Dict[str, str]) -> List[str]: - cmd = [] - if "ContentType" in extra_args: - cmd += ["--content-type", extra_args["ContentType"]] - if "ContentEncoding" in extra_args: - cmd += ["--content-encoding", extra_args["ContentEncoding"]] - if "ACL" in extra_args: - cmd += ["--acl", extra_args["ACL"]] - return cmd - - -class AwsS3Proxy(_common_data.DataProxy): - _AWS_CLI = "aws" - _SHARD_CHARACTERS = [_text_type(x) for x in _six_moves.range(10)] + list(_string.ascii_lowercase) - - def __init__(self, raw_output_data_prefix_override: str = None): - """ - :param raw_output_data_prefix_override: Instead of relying on the AWS or GCS configuration (see - S3_SHARD_FORMATTER for AWS and GCS_PREFIX for GCP) setting when computing the shard - path (_get_shard_path), use this prefix instead as a base. This code assumes that the - path passed in is correct. That is, an S3 path won't be passed in when running on GCP. - """ - super().__init__(name="awscli-s3") - self._raw_output_data_prefix_override = raw_output_data_prefix_override - - @property - def raw_output_data_prefix_override(self) -> str: - return self._raw_output_data_prefix_override - - @staticmethod - def _check_binary(): - """ - Make sure that the AWS cli is present - """ - if not _which(AwsS3Proxy._AWS_CLI): - raise _FlyteUserException("AWS CLI not found at Please install.") - - @staticmethod - def _split_s3_path_to_bucket_and_key(path): - """ - :param Text path: - :rtype: (Text, Text) - """ - path = path[len("s3://") :] - first_slash = path.index("/") - return path[:first_slash], path[first_slash + 1 :] - - def exists(self, remote_path): - """ - :param Text remote_path: remote s3:// path - :rtype bool: whether the s3 file exists or not - """ - AwsS3Proxy._check_binary() - - if not remote_path.startswith("s3://"): - raise ValueError("Not an S3 ARN. Please use FQN (S3 ARN) of the format s3://...") - - bucket, file_path = self._split_s3_path_to_bucket_and_key(remote_path) - cmd = [ - AwsS3Proxy._AWS_CLI, - "s3api", - "head-object", - "--bucket", - bucket, - "--key", - file_path, - ] - try: - _update_cmd_config_and_execute(cmd) - return True - except Exception as ex: - # The s3api command returns an error if the object does not exist. The error message contains - # the http status code: "An error occurred (404) when calling the HeadObject operation: Not Found" - # This is a best effort for returning if the object does not exist by searching - # for existence of (404) in the error message. This should not be needed when we get off the cli and use lib - if _re.search("(404)", _text_type(ex)): - return False - else: - raise ex - - def download_directory(self, remote_path, local_path): - """ - :param Text remote_path: remote s3:// path - :param Text local_path: directory to copy to - """ - AwsS3Proxy._check_binary() - - if not remote_path.startswith("s3://"): - raise ValueError("Not an S3 ARN. Please use FQN (S3 ARN) of the format s3://...") - - cmd = [AwsS3Proxy._AWS_CLI, "s3", "cp", "--recursive", remote_path, local_path] - return _update_cmd_config_and_execute(cmd) - - def download(self, remote_path, local_path): - """ - :param Text remote_path: remote s3:// path - :param Text local_path: directory to copy to - """ - if not remote_path.startswith("s3://"): - raise ValueError("Not an S3 ARN. Please use FQN (S3 ARN) of the format s3://...") - - AwsS3Proxy._check_binary() - cmd = [AwsS3Proxy._AWS_CLI, "s3", "cp", remote_path, local_path] - return _update_cmd_config_and_execute(cmd) - - def upload(self, file_path, to_path): - """ - :param Text file_path: - :param Text to_path: - """ - AwsS3Proxy._check_binary() - - extra_args = { - "ACL": "bucket-owner-full-control", - } - - cmd = [AwsS3Proxy._AWS_CLI, "s3", "cp"] - cmd.extend(_extra_args(extra_args)) - cmd += [file_path, to_path] - - return _update_cmd_config_and_execute(cmd) - - def upload_directory(self, local_path, remote_path): - """ - :param Text local_path: - :param Text remote_path: - """ - extra_args = { - "ACL": "bucket-owner-full-control", - } - - if not remote_path.startswith("s3://"): - raise ValueError("Not an S3 ARN. Please use FQN (S3 ARN) of the format s3://...") - - AwsS3Proxy._check_binary() - cmd = [AwsS3Proxy._AWS_CLI, "s3", "cp", "--recursive"] - cmd.extend(_extra_args(extra_args)) - cmd += [local_path, remote_path] - return _update_cmd_config_and_execute(cmd) - - def get_random_path(self): - """ - :rtype: Text - """ - # Create a 128-bit random hash because the birthday attack principle shows that there is about a 50% chance of a - # collision between objects when 2^(n/2) objects are created (where n is the number of bits in the hash). - # Assuming Flyte eventually creates 1 trillion pieces of data (~2 ^ 40), the likelihood - # of a collision is 10^-15 with 128-bit...or basically 0. - key = _uuid.UUID(int=_flyte_random.random.getrandbits(128)).hex - return _os.path.join(self._get_shard_path(), key) - - def get_random_directory(self): - """ - :rtype: Text - """ - return self.get_random_path() + "/" - - def _get_shard_path(self) -> str: - """ - If this object was created with a raw output data prefix, usually set by Propeller/Plugins at execution time - and piped all the way here, it will be used instead of referencing the S3 shard configuration. - """ - if self.raw_output_data_prefix_override: - return self.raw_output_data_prefix_override - - shard = "" - for _ in _six_moves.range(_aws_config.S3_SHARD_STRING_LENGTH.get()): - shard += _flyte_random.random.choice(self._SHARD_CHARACTERS) - return _aws_config.S3_SHARD_FORMATTER.get().format(shard) diff --git a/flytekit/models/common.py b/flytekit/models/common.py index 3ef6291ac7..63253bf399 100644 --- a/flytekit/models/common.py +++ b/flytekit/models/common.py @@ -1,7 +1,6 @@ import abc as _abc import json as _json -import six as _six from flyteidl.admin import common_pb2 as _common_pb2 from google.protobuf import json_format as _json_format from google.protobuf import struct_pb2 as _struct @@ -58,7 +57,7 @@ def short_string(self): """ :rtype: Text """ - return _six.text_type(self.to_flyte_idl()) + return str(self.to_flyte_idl()) def verbose_string(self): """ @@ -333,7 +332,7 @@ def to_flyte_idl(self): """ :rtype: dict[Text, Text] """ - return _common_pb2.Labels(values={k: v for k, v in _six.iteritems(self.values)}) + return _common_pb2.Labels(values={k: v for k, v in self.values.items()}) @classmethod def from_flyte_idl(cls, pb2_object): @@ -341,7 +340,7 @@ def from_flyte_idl(cls, pb2_object): :param flyteidl.admin.common_pb2.Labels pb2_object: :rtype: Labels """ - return cls({k: v for k, v in _six.iteritems(pb2_object.values)}) + return cls({k: v for k, v in pb2_object.values.items()}) class Annotations(FlyteIdlEntity): @@ -361,7 +360,7 @@ def to_flyte_idl(self): """ :rtype: _common_pb2.Annotations """ - return _common_pb2.Annotations(values={k: v for k, v in _six.iteritems(self.values)}) + return _common_pb2.Annotations(values={k: v for k, v in self.values.items()}) @classmethod def from_flyte_idl(cls, pb2_object): @@ -369,7 +368,7 @@ def from_flyte_idl(cls, pb2_object): :param flyteidl.admin.common_pb2.Annotations pb2_object: :rtype: Annotations """ - return cls({k: v for k, v in _six.iteritems(pb2_object.values)}) + return cls({k: v for k, v in pb2_object.values.items()}) class UrlBlob(FlyteIdlEntity): diff --git a/flytekit/models/core/compiler.py b/flytekit/models/core/compiler.py index 3246ee22b3..929f816227 100644 --- a/flytekit/models/core/compiler.py +++ b/flytekit/models/core/compiler.py @@ -1,4 +1,3 @@ -import six as _six from flyteidl.core import compiler_pb2 as _compiler_pb2 from flytekit.models import common as _common @@ -61,8 +60,8 @@ def to_flyte_idl(self): :rtype: flyteidl.core.compiler_pb2.ConnectionSet """ return _compiler_pb2.ConnectionSet( - upstream={k: v.to_flyte_idl() for k, v in _six.iteritems(self.upstream)}, - downstream={k: v.to_flyte_idl() for k, v in _six.iteritems(self.upstream)}, + upstream={k: v.to_flyte_idl() for k, v in self.upstream.items()}, + downstream={k: v.to_flyte_idl() for k, v in self.upstream.items()}, ) @classmethod @@ -72,8 +71,8 @@ def from_flyte_idl(cls, p): :rtype: ConnectionSet """ return cls( - upstream={k: ConnectionSet.IdList.from_flyte_idl(v) for k, v in _six.iteritems(p.upstream)}, - downstream={k: ConnectionSet.IdList.from_flyte_idl(v) for k, v in _six.iteritems(p.downstream)}, + upstream={k: ConnectionSet.IdList.from_flyte_idl(v) for k, v in p.upstream.items()}, + downstream={k: ConnectionSet.IdList.from_flyte_idl(v) for k, v in p.downstream.items()}, ) diff --git a/flytekit/models/interface.py b/flytekit/models/interface.py index 5364d39c3e..b0c9fc882a 100644 --- a/flytekit/models/interface.py +++ b/flytekit/models/interface.py @@ -1,6 +1,5 @@ import typing -import six as _six from flyteidl.core import interface_pb2 as _interface_pb2 from flytekit.models import common as _common @@ -73,7 +72,7 @@ def to_flyte_idl(self): """ :rtype: dict[Text, Variable] """ - return _interface_pb2.VariableMap(variables={k: v.to_flyte_idl() for k, v in _six.iteritems(self.variables)}) + return _interface_pb2.VariableMap(variables={k: v.to_flyte_idl() for k, v in self.variables.items()}) @classmethod def from_flyte_idl(cls, pb2_object): @@ -81,7 +80,7 @@ def from_flyte_idl(cls, pb2_object): :param dict[Text, Variable] pb2_object: :rtype: VariableMap """ - return cls({k: Variable.from_flyte_idl(v) for k, v in _six.iteritems(pb2_object.variables)}) + return cls({k: Variable.from_flyte_idl(v) for k, v in pb2_object.variables.items()}) class TypedInterface(_common.FlyteIdlEntity): @@ -106,10 +105,8 @@ def outputs(self) -> typing.Dict[str, Variable]: def to_flyte_idl(self) -> _interface_pb2.TypedInterface: return _interface_pb2.TypedInterface( - inputs=_interface_pb2.VariableMap(variables={k: v.to_flyte_idl() for k, v in _six.iteritems(self.inputs)}), - outputs=_interface_pb2.VariableMap( - variables={k: v.to_flyte_idl() for k, v in _six.iteritems(self.outputs)} - ), + inputs=_interface_pb2.VariableMap(variables={k: v.to_flyte_idl() for k, v in self.inputs.items()}), + outputs=_interface_pb2.VariableMap(variables={k: v.to_flyte_idl() for k, v in self.outputs.items()}), ) @classmethod @@ -118,8 +115,8 @@ def from_flyte_idl(cls, proto: _interface_pb2.TypedInterface) -> "TypedInterface :param proto: """ return cls( - inputs={k: Variable.from_flyte_idl(v) for k, v in _six.iteritems(proto.inputs.variables)}, - outputs={k: Variable.from_flyte_idl(v) for k, v in _six.iteritems(proto.outputs.variables)}, + inputs={k: Variable.from_flyte_idl(v) for k, v in proto.inputs.variables.items()}, + outputs={k: Variable.from_flyte_idl(v) for k, v in proto.outputs.variables.items()}, ) @@ -211,7 +208,7 @@ def to_flyte_idl(self): :rtype: flyteidl.core.interface_pb2.ParameterMap """ return _interface_pb2.ParameterMap( - parameters={k: v.to_flyte_idl() for k, v in _six.iteritems(self.parameters)}, + parameters={k: v.to_flyte_idl() for k, v in self.parameters.items()}, ) @classmethod @@ -220,4 +217,4 @@ def from_flyte_idl(cls, pb2_object): :param flyteidl.core.interface_pb2.ParameterMap pb2_object: :rtype: ParameterMap """ - return cls(parameters={k: Parameter.from_flyte_idl(v) for k, v in _six.iteritems(pb2_object.parameters)}) + return cls(parameters={k: Parameter.from_flyte_idl(v) for k, v in pb2_object.parameters.items()}) diff --git a/flytekit/models/literals.py b/flytekit/models/literals.py index 9333f38ba2..6584334450 100644 --- a/flytekit/models/literals.py +++ b/flytekit/models/literals.py @@ -5,7 +5,7 @@ from flyteidl.core import literals_pb2 as _literals_pb2 from google.protobuf.struct_pb2 import Struct -from flytekit.common.exceptions import user as _user_exceptions +from flytekit.exceptions import user as _user_exceptions from flytekit.models import common as _common from flytekit.models.core import types as _core_types from flytekit.models.types import LiteralType as _LiteralType diff --git a/flytekit/models/task.py b/flytekit/models/task.py index 76caa0444d..caf2471e58 100644 --- a/flytekit/models/task.py +++ b/flytekit/models/task.py @@ -1,23 +1,18 @@ import json as _json import typing -import six as _six from flyteidl.admin import task_pb2 as _admin_task from flyteidl.core import compiler_pb2 as _compiler from flyteidl.core import literals_pb2 as _literals_pb2 from flyteidl.core import tasks_pb2 as _core_task -from flyteidl.plugins import spark_pb2 as _spark_task from google.protobuf import json_format as _json_format from google.protobuf import struct_pb2 as _struct -from flytekit.common.exceptions import user as _user_exceptions from flytekit.models import common as _common from flytekit.models import interface as _interface from flytekit.models import literals as _literals from flytekit.models import security as _sec from flytekit.models.core import identifier as _identifier -from flytekit.plugins import flyteidl as _lazy_flyteidl -from flytekit.sdk.spark_types import SparkType as _spark_type class Resources(_common.FlyteIdlEntity): @@ -617,151 +612,6 @@ def from_flyte_idl(cls, pb2_object): return cls(template=TaskTemplate.from_flyte_idl(pb2_object.template)) -class SparkJob(_common.FlyteIdlEntity): - """ - This model is deprecated and will be removed in 1.0.0. Please use the definition in the - flytekit spark plugin instead. - """ - - def __init__( - self, - spark_type, - application_file, - main_class, - spark_conf, - hadoop_conf, - executor_path, - ): - """ - This defines a SparkJob target. It will execute the appropriate SparkJob. - - :param application_file: The main application file to execute. - :param dict[Text, Text] spark_conf: A definition of key-value pairs for spark config for the job. - :param dict[Text, Text] hadoop_conf: A definition of key-value pairs for hadoop config for the job. - """ - self._application_file = application_file - self._spark_type = spark_type - self._main_class = main_class - self._executor_path = executor_path - self._spark_conf = spark_conf - self._hadoop_conf = hadoop_conf - - def with_overrides( - self, new_spark_conf: typing.Dict[str, str] = None, new_hadoop_conf: typing.Dict[str, str] = None - ) -> "SparkJob": - if not new_spark_conf: - new_spark_conf = self.spark_conf - - if not new_hadoop_conf: - new_hadoop_conf = self.hadoop_conf - - return SparkJob( - spark_type=self.spark_type, - application_file=self.application_file, - main_class=self.main_class, - spark_conf=new_spark_conf, - hadoop_conf=new_hadoop_conf, - executor_path=self.executor_path, - ) - - @property - def main_class(self): - """ - The main class to execute - :rtype: Text - """ - return self._main_class - - @property - def spark_type(self): - """ - Spark Job Type - :rtype: Text - """ - return self._spark_type - - @property - def application_file(self): - """ - The main application file to execute - :rtype: Text - """ - return self._application_file - - @property - def executor_path(self): - """ - The python executable to use - :rtype: Text - """ - return self._executor_path - - @property - def spark_conf(self): - """ - A definition of key-value pairs for spark config for the job. - :rtype: dict[Text, Text] - """ - return self._spark_conf - - @property - def hadoop_conf(self): - """ - A definition of key-value pairs for hadoop config for the job. - :rtype: dict[Text, Text] - """ - return self._hadoop_conf - - def to_flyte_idl(self): - """ - :rtype: flyteidl.plugins.spark_pb2.SparkJob - """ - - if self.spark_type == _spark_type.PYTHON: - application_type = _spark_task.SparkApplication.PYTHON - elif self.spark_type == _spark_type.JAVA: - application_type = _spark_task.SparkApplication.JAVA - elif self.spark_type == _spark_type.SCALA: - application_type = _spark_task.SparkApplication.SCALA - elif self.spark_type == _spark_type.R: - application_type = _spark_task.SparkApplication.R - else: - raise _user_exceptions.FlyteValidationException("Invalid Spark Application Type Specified") - - return _spark_task.SparkJob( - applicationType=application_type, - mainApplicationFile=self.application_file, - mainClass=self.main_class, - executorPath=self.executor_path, - sparkConf=self.spark_conf, - hadoopConf=self.hadoop_conf, - ) - - @classmethod - def from_flyte_idl(cls, pb2_object): - """ - :param flyteidl.plugins.spark_pb2.SparkJob pb2_object: - :rtype: SparkJob - """ - - application_type = _spark_type.PYTHON - if pb2_object.type == _spark_task.SparkApplication.JAVA: - application_type = _spark_type.JAVA - elif pb2_object.type == _spark_task.SparkApplication.SCALA: - application_type = _spark_type.SCALA - elif pb2_object.type == _spark_task.SparkApplication.R: - application_type = _spark_type.R - - return cls( - type=application_type, - spark_conf=pb2_object.sparkConf, - application_file=pb2_object.mainApplicationFile, - main_class=pb2_object.mainClass, - hadoop_conf=pb2_object.hadoopConf, - executor_path=pb2_object.executorPath, - ) - - class IOStrategy(_common.FlyteIdlEntity): """ Provides methods to manage data in and out of the Raw container using Download Modes. This can only be used if DataLoadingConfig is enabled. @@ -930,8 +780,8 @@ def to_flyte_idl(self): command=self.command, args=self.args, resources=self.resources.to_flyte_idl(), - env=[_literals_pb2.KeyValuePair(key=k, value=v) for k, v in _six.iteritems(self.env)], - config=[_literals_pb2.KeyValuePair(key=k, value=v) for k, v in _six.iteritems(self.config)], + env=[_literals_pb2.KeyValuePair(key=k, value=v) for k, v in self.env.items()], + config=[_literals_pb2.KeyValuePair(key=k, value=v) for k, v in self.config.items()], data_config=self._data_loading_config.to_flyte_idl() if self._data_loading_config else None, ) @@ -1045,72 +895,3 @@ def from_flyte_idl(cls, pb2_object: _core_task.Sql): statement=pb2_object.statement, dialect=pb2_object.dialect, ) - - -class SidecarJob(_common.FlyteIdlEntity): - def __init__(self, pod_spec, primary_container_name, annotations=None, labels=None): - """ - A sidecar job represents the full kubernetes pod spec and related metadata required for executing a sidecar - task. - - :param pod_spec: k8s.io.api.core.v1.PodSpec - :param primary_container_name: Text - :param dict[Text, Text] annotations: - :param dict[Text, Text] labels: - """ - self._pod_spec = pod_spec - self._primary_container_name = primary_container_name - self._annotations = annotations - self._labels = labels - - @property - def pod_spec(self): - """ - :rtype: k8s.io.api.core.v1.PodSpec - """ - return self._pod_spec - - @property - def primary_container_name(self): - """ - :rtype: Text - """ - return self._primary_container_name - - @property - def annotations(self): - """ - :rtype: dict[Text,Text] - """ - return self._annotations - - @property - def labels(self): - """ - :rtype: dict[Text,Text] - """ - return self._labels - - def to_flyte_idl(self): - """ - :rtype: flyteidl.core.tasks_pb2.SidecarJob - """ - return _lazy_flyteidl.plugins.sidecar_pb2.SidecarJob( - pod_spec=self.pod_spec, - primary_container_name=self.primary_container_name, - annotations=self.annotations, - labels=self.labels, - ) - - @classmethod - def from_flyte_idl(cls, pb2_object): - """ - :param flyteidl.admin.task_pb2.Task pb2_object: - :rtype: Container - """ - return cls( - pod_spec=pb2_object.pod_spec, - primary_container_name=pb2_object.primary_container_name, - annotations=pb2_object.annotations, - labels=pb2_object.labels, - ) diff --git a/flytekit/plugins/__init__.py b/flytekit/plugins/__init__.py deleted file mode 100644 index 61333933f1..0000000000 --- a/flytekit/plugins/__init__.py +++ /dev/null @@ -1,34 +0,0 @@ -""" -This file is for old style plugins - for new plugins that work with the Python native-typed Flytekit, please -refer to the plugin specific directory underneath the plugins folder at the top level of this repository. -""" -from flytekit.tools import lazy_loader as _lazy_loader - -pyspark = _lazy_loader.lazy_load_module("pyspark") # type: _lazy_loader._LazyLoadModule - -k8s = _lazy_loader.lazy_load_module("k8s") # type: _lazy_loader._LazyLoadModule -type(k8s).add_sub_module("io.api.core.v1.generated_pb2") -type(k8s).add_sub_module("io.apimachinery.pkg.api.resource.generated_pb2") - -flyteidl = _lazy_loader.lazy_load_module("flyteidl") # type: _lazy_loader._LazyLoadModule -type(flyteidl).add_sub_module("plugins.sidecar_pb2") - -numpy = _lazy_loader.lazy_load_module("numpy") # type: _lazy_loader._LazyLoadModule -pandas = _lazy_loader.lazy_load_module("pandas") # type: _lazy_loader._LazyLoadModule - -hmsclient = _lazy_loader.lazy_load_module("hmsclient") # type: _lazy_loader._LazyLoadModule -type(hmsclient).add_sub_module("genthrift.hive_metastore.ttypes") - -_lazy_loader.LazyLoadPlugin("spark", ["pyspark>=2.4.0,<3.0.0"], [pyspark]) - -_lazy_loader.LazyLoadPlugin("spark3", ["pyspark>=3.0.0"], [pyspark]) - -_lazy_loader.LazyLoadPlugin("sidecar", ["k8s-proto>=0.0.3,<1.0.0"], [k8s, flyteidl]) - -_lazy_loader.LazyLoadPlugin( - "schema", - ["numpy>=1.14.0,<2.0.0", "pandas>=0.22.0,<2.0.0", "pyarrow>=0.11.0,<1.0.0"], - [numpy, pandas], -) - -_lazy_loader.LazyLoadPlugin("hive_sensor", ["hmsclient>=0.0.1,<1.0.0"], [hmsclient]) diff --git a/flytekit/remote/component_nodes.py b/flytekit/remote/component_nodes.py index 367cab8997..877a6d6494 100644 --- a/flytekit/remote/component_nodes.py +++ b/flytekit/remote/component_nodes.py @@ -1,7 +1,7 @@ import logging as _logging from typing import Dict -from flytekit.common.exceptions import system as _system_exceptions +from flytekit.exceptions import system as _system_exceptions from flytekit.models import launch_plan as _launch_plan_model from flytekit.models import task as _task_model from flytekit.models.core import identifier as id_models diff --git a/flytekit/remote/executions.py b/flytekit/remote/executions.py index b581769f52..64b3fcb7d8 100644 --- a/flytekit/remote/executions.py +++ b/flytekit/remote/executions.py @@ -2,9 +2,9 @@ from typing import Any, Dict, List, Optional, Union -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.exceptions import user as user_exceptions from flytekit.core.type_engine import LiteralsResolver +from flytekit.exceptions import user as _user_exceptions +from flytekit.exceptions import user as user_exceptions from flytekit.models import execution as execution_models from flytekit.models import node_execution as node_execution_models from flytekit.models.admin import task_execution as admin_task_execution_models diff --git a/flytekit/remote/launch_plan.py b/flytekit/remote/launch_plan.py index 016e3a3489..b5dddea03f 100644 --- a/flytekit/remote/launch_plan.py +++ b/flytekit/remote/launch_plan.py @@ -1,10 +1,7 @@ from typing import Optional -from flytekit.common.exceptions import scopes as _exception_scopes -from flytekit.common.exceptions import user as _user_exceptions from flytekit.core.interface import Interface from flytekit.core.type_engine import TypeEngine -from flytekit.engines.flyte import engine as _flyte_engine from flytekit.models import interface as _interface_models from flytekit.models import launch_plan as _launch_plan_models from flytekit.models.core import identifier as id_models @@ -94,14 +91,5 @@ def guessed_python_interface(self, value): return self._python_interface = value - @_exception_scopes.system_entry_point - def update(self, state: _launch_plan_models.LaunchPlanState): - if not self.id: - raise _user_exceptions.FlyteAssertion( - "Failed to update launch plan because the launch plan's ID is not set. Please call register to fetch " - "or register the identifier first" - ) - return _flyte_engine.get_client().update_launch_plan(self.id, state) - def __repr__(self) -> str: return f"FlyteLaunchPlan(ID: {self.id} Interface: {self.interface} WF ID: {self.workflow_id})" diff --git a/flytekit/remote/nodes.py b/flytekit/remote/nodes.py index f8ae1b2d6a..4131905c5b 100644 --- a/flytekit/remote/nodes.py +++ b/flytekit/remote/nodes.py @@ -3,11 +3,11 @@ import logging as _logging from typing import Dict, List, Optional, Union -from flytekit.common import constants as _constants -from flytekit.common.exceptions import system as _system_exceptions -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.mixins import hash as _hash_mixin +from flytekit.core import constants as _constants +from flytekit.core import hash as _hash_mixin from flytekit.core.promise import NodeOutput +from flytekit.exceptions import system as _system_exceptions +from flytekit.exceptions import user as _user_exceptions from flytekit.models import launch_plan as _launch_plan_model from flytekit.models import task as _task_model from flytekit.models.core import identifier as id_models diff --git a/flytekit/remote/remote.py b/flytekit/remote/remote.py index 2178f148d3..44d5e8617c 100644 --- a/flytekit/remote/remote.py +++ b/flytekit/remote/remote.py @@ -19,14 +19,14 @@ from flyteidl.core import literals_pb2 as literals_pb2 from flytekit.clients.friendly import SynchronousFlyteClient -from flytekit.common import utils as common_utils -from flytekit.common.exceptions.user import FlyteEntityAlreadyExistsException, FlyteEntityNotExistException from flytekit.configuration import internal from flytekit.configuration import platform as platform_config from flytekit.configuration import sdk as sdk_config from flytekit.configuration import set_flyte_config_file -from flytekit.core import context_manager +from flytekit.core import constants, context_manager, utils from flytekit.core.interface import Interface +from flytekit.exceptions import user as user_exceptions +from flytekit.exceptions.user import FlyteEntityAlreadyExistsException, FlyteEntityNotExistException from flytekit.loggers import remote_logger from flytekit.models import filters as filter_models from flytekit.models.admin import common as admin_common_models @@ -39,9 +39,6 @@ from flytekit.clients.helpers import iterate_node_executions, iterate_task_executions from flytekit.clis.flyte_cli.main import _detect_default_config_file from flytekit.clis.sdk_in_container import serialize -from flytekit.common import constants -from flytekit.common.exceptions import user as user_exceptions -from flytekit.common.translator import FlyteControlPlaneEntity, FlyteLocalEntity, get_serializable from flytekit.configuration import auth as auth_config from flytekit.configuration.internal import DOMAIN, PROJECT from flytekit.core.base_task import PythonTask @@ -68,6 +65,7 @@ from flytekit.remote.nodes import FlyteNode from flytekit.remote.task import FlyteTask from flytekit.remote.workflow import FlyteWorkflow +from flytekit.tools.translator import FlyteControlPlaneEntity, FlyteLocalEntity, get_serializable ExecutionDataResponse = typing.Union[WorkflowExecutionGetDataResponse, NodeExecutionGetDataResponse] @@ -1297,7 +1295,7 @@ def _get_input_literal_map(self, execution_data: ExecutionDataResponse) -> liter tmp_name = os.path.join(ctx.file_access.local_sandbox_dir, "inputs.pb") ctx.file_access.get_data(execution_data.inputs.url, tmp_name) return literal_models.LiteralMap.from_flyte_idl( - common_utils.load_proto_from_file(literals_pb2.LiteralMap, tmp_name) + utils.load_proto_from_file(literals_pb2.LiteralMap, tmp_name) ) return literal_models.LiteralMap({}) @@ -1310,6 +1308,6 @@ def _get_output_literal_map(self, execution_data: ExecutionDataResponse) -> lite tmp_name = os.path.join(ctx.file_access.local_sandbox_dir, "outputs.pb") ctx.file_access.get_data(execution_data.outputs.url, tmp_name) return literal_models.LiteralMap.from_flyte_idl( - common_utils.load_proto_from_file(literals_pb2.LiteralMap, tmp_name) + utils.load_proto_from_file(literals_pb2.LiteralMap, tmp_name) ) return literal_models.LiteralMap({}) diff --git a/flytekit/remote/task.py b/flytekit/remote/task.py index 0c48f5f15e..1ff99549d6 100644 --- a/flytekit/remote/task.py +++ b/flytekit/remote/task.py @@ -1,6 +1,6 @@ from typing import Optional -from flytekit.common.mixins import hash as _hash_mixin +from flytekit.core import hash as _hash_mixin from flytekit.core.interface import Interface from flytekit.core.type_engine import TypeEngine from flytekit.loggers import logger diff --git a/flytekit/remote/workflow.py b/flytekit/remote/workflow.py index 396f377500..010a703187 100644 --- a/flytekit/remote/workflow.py +++ b/flytekit/remote/workflow.py @@ -2,11 +2,11 @@ from typing import Dict, List, Optional -from flytekit.common import constants as _constants -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.mixins import hash as _hash_mixin +from flytekit.core import constants as _constants +from flytekit.core import hash as _hash_mixin from flytekit.core.interface import Interface from flytekit.core.type_engine import TypeEngine +from flytekit.exceptions import user as _user_exceptions from flytekit.models import launch_plan as launch_plan_models from flytekit.models import task as _task_models from flytekit.models.core import compiler as compiler_models diff --git a/flytekit/sdk/__init__.py b/flytekit/sdk/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/flytekit/sdk/exceptions.py b/flytekit/sdk/exceptions.py deleted file mode 100644 index 89874a5518..0000000000 --- a/flytekit/sdk/exceptions.py +++ /dev/null @@ -1,11 +0,0 @@ -from flytekit.common.exceptions import user as _user - - -class RecoverableException(_user.FlyteRecoverableException): - """ - Raise an exception of this type if user code detects an error and would like to force a retry of the entire task. - Any exception raised from user code other than RecoverableException will NOT be considered retryable and the task - will fail without additional retries. - """ - - pass diff --git a/flytekit/sdk/spark_types.py b/flytekit/sdk/spark_types.py deleted file mode 100644 index 9477895fac..0000000000 --- a/flytekit/sdk/spark_types.py +++ /dev/null @@ -1,8 +0,0 @@ -import enum - - -class SparkType(enum.Enum): - PYTHON = 1 - SCALA = 2 - JAVA = 3 - R = 4 diff --git a/flytekit/sdk/tasks.py b/flytekit/sdk/tasks.py deleted file mode 100644 index 0c73ad66b7..0000000000 --- a/flytekit/sdk/tasks.py +++ /dev/null @@ -1,1244 +0,0 @@ -import datetime as _datetime - -import six as _six - -from flytekit.common import constants as _common_constants -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.tasks import generic_spark_task as _sdk_generic_spark_task -from flytekit.common.tasks import hive_task as _sdk_hive_tasks -from flytekit.common.tasks import sdk_dynamic as _sdk_dynamic -from flytekit.common.tasks import sdk_runnable as _sdk_runnable_tasks -from flytekit.common.tasks import sidecar_task as _sdk_sidecar_tasks -from flytekit.common.tasks import spark_task as _sdk_spark_tasks -from flytekit.common.tasks import task as _task -from flytekit.common.types import helpers as _type_helpers -from flytekit.models import interface as _interface_model -from flytekit.sdk.spark_types import SparkType as _spark_type - - -def inputs(_task_template=None, **kwargs): - """ - Decorator that provides input definitions to a decorated task definition. - - .. note:: - - Certain tasks have special input behavior. See comments on each task decorator for more information. - - .. code-block:: python - - @inputs(in1=Types.Integer, in2=[Types.String], in3=[[[Types.Integer]]]) - @outputs(out1=Types.Integer, out2=Types.String) - @python_task - def my_task(wf_params, in1, in2, out1, out2): - pass - - :param flytekit.common.tasks.sdk_runnable.SdkRunnableTask _task_template: Do not declare directly. This is the - decorated task template. - :param dict[Text,flytekit.common.types.base_sdk_types.FlyteSdkType] kwargs: Arbitrary keyword arguments for input - name and type. - :rtype: flytekit.common.tasks.sdk_runnable.SdkRunnableTask - """ - - def apply_inputs_wrapper(task): - if not isinstance(task, _task.SdkTask): - additional_msg = ( - "Inputs can only be applied to a task. Did you forget the task decorator on method '{}.{}'?".format( - task.__module__, - task.__name__ if hasattr(task, "__name__") else "", - ) - ) - raise _user_exceptions.FlyteTypeException( - expected_type=_sdk_runnable_tasks.SdkRunnableTask, - received_type=type(task), - received_value=task, - additional_msg=additional_msg, - ) - for k, v in _six.iteritems(kwargs): - kwargs[k] = _interface_model.Variable( - _type_helpers.python_std_to_sdk_type(v).to_flyte_literal_type(), "" - ) # TODO: Support descriptions - - task.add_inputs(kwargs) - return task - - if _task_template is not None: - return apply_inputs_wrapper(_task_template) - else: - return apply_inputs_wrapper - - -def outputs(_task_template=None, **kwargs): - """ - Decorator that provides output definitions to a decorated task definition. - - .. note:: - - Certain tasks have special output behavior. See comments on each task decorator for more information. - - .. code-block:: python - - @outputs(out1=Types.Integer, out2=Types.String) - @python_task - def my_task(wf_params, out1, out2): - out1.set(123) - out2.set('hello world!') - - :param flytekit.common.tasks.sdk_runnable.SdkRunnableTask _task_template: Do not declare directly. This is the - decorated task template. - :param dict[Text,flytekit.common.types.base_sdk_types.FlyteSdkType] kwargs: Arbitrary keyword arguments for input - name and type. - :rtype: flytekit.common.tasks.sdk_runnable.SdkRunnableTask - """ - - def apply_outputs_wrapper(task): - if not isinstance(task, _sdk_runnable_tasks.SdkRunnableTask) and not isinstance( - task, _nb_tasks.SdkNotebookTask - ): - additional_msg = ( - "Outputs can only be applied to a task. Did you forget the task decorator on method '{}.{}'?".format( - task.__module__, - task.__name__ if hasattr(task, "__name__") else "", - ) - ) - raise _user_exceptions.FlyteTypeException( - expected_type=_sdk_runnable_tasks.SdkRunnableTask, - received_type=type(task), - received_value=task, - additional_msg=additional_msg, - ) - for k, v in _six.iteritems(kwargs): - kwargs[k] = _interface_model.Variable( - _type_helpers.python_std_to_sdk_type(v).to_flyte_literal_type(), "" - ) # TODO: Support descriptions - - task.add_outputs(kwargs) - return task - - if _task_template is not None: - return apply_outputs_wrapper(_task_template) - else: - return apply_outputs_wrapper - - -def python_task( - _task_function=None, - cache_version="", - retries=0, - interruptible=None, - deprecated="", - storage_request=None, - cpu_request=None, - gpu_request=None, - memory_request=None, - storage_limit=None, - cpu_limit=None, - gpu_limit=None, - memory_limit=None, - cache=False, - timeout=None, - environment=None, - cache_serialize=False, - cls=None, -): - """ - Decorator to create a Python Task definition. This task will run as a single unit of work on the platform. - - .. code-block:: python - - @inputs(int_list=[Types.Integer]) - @outputs(sum_of_list=Types.Integer - @python_task - def my_task(wf_params, int_list, sum_of_list): - sum_of_list.set(sum(int_list)) - - :param _task_function: this is the decorated method and shouldn't be declared explicitly. The function must - take a first argument, and then named arguments matching those defined in @inputs and @outputs. No keyword - arguments are allowed for wrapped task functions. - - :param Text cache_version: [optional] string representing logical version for discovery. This field should be - updated whenever the underlying algorithm changes. - - .. note:: - - This argument is required to be a non-empty string if `cache` is True. - - :param int retries: [optional] integer determining number of times task can be retried on - :py:exc:`flytekit.sdk.exceptions.RecoverableException` or transient platform failures. Defaults - to 0. - - .. note:: - - If retries > 0, the task must be able to recover from any remote state created within the user code. It is - strongly recommended that tasks are written to be idempotent. - - :param bool interruptible: [optional] boolean describing if the task is interruptible. - - :param Text deprecated: [optional] string that should be provided if this task is deprecated. The string - will be logged as a warning so it should contain information regarding how to update to a newer task. - - :param Text storage_request: [optional] Kubernetes resource string for lower-bound of disk storage space - for the task to run. Default is set by platform-level configuration. - - .. note:: - - This is currently not supported by the platform. - - :param Text cpu_request: [optional] Kubernetes resource string for lower-bound of cores for the task to execute. - This can be set to a fractional portion of a CPU. Default is set by platform-level configuration. - - TODO: Add links to resource string documentation for Kubernetes - - :param Text gpu_request: [optional] Kubernetes resource string for lower-bound of desired GPUs. - Default is set by platform-level configuration. - - TODO: Add links to resource string documentation for Kubernetes - - :param Text memory_request: [optional] Kubernetes resource string for lower-bound of physical memory - necessary for the task to execute. Default is set by platform-level configuration. - - TODO: Add links to resource string documentation for Kubernetes - - :param Text storage_limit: [optional] Kubernetes resource string for upper-bound of disk storage space - for the task to run. This amount is not guaranteed! If not specified, it is set equal to storage_request. - - .. note:: - - This is currently not supported by the platform. - - :param Text cpu_limit: [optional] Kubernetes resource string for upper-bound of cores for the task to execute. - This can be set to a fractional portion of a CPU. This amount is not guaranteed! If not specified, - it is set equal to cpu_request. - - :param Text gpu_limit: [optional] Kubernetes resource string for upper-bound of desired GPUs. This amount is not - guaranteed! If not specified, it is set equal to gpu_request. - - :param Text memory_limit: [optional] Kubernetes resource string for upper-bound of physical memory - necessary for the task to execute. This amount is not guaranteed! If not specified, it is set equal to - memory_request. - - :param bool cache: [optional] boolean describing if the outputs of this task should be cached and - re-usable. - - :param datetime.timedelta timeout: [optional] describes how long the task should be allowed to - run at max before triggering a retry (if retries are enabled). By default, tasks are allowed to run - indefinitely. If a null timedelta is passed (i.e. timedelta(seconds=0)), the task will not timeout. - - :param dict[Text,Text] environment: [optional] environment variables to set when executing this task. - - :param bool cache_serialize: [optional] boolean describing if instances of this cachable task should be executed - in serial. This means only a single instances executes and other concurrent executions wait for it to complete - and reuse the cached outputs. - - :param cls: This can be used to override the task implementation with a user-defined extension. The class - provided must be a subclass of flytekit.common.tasks.sdk_runnable.SdkRunnableTask. A user can use this to - inject bespoke logic into the base Flyte programming model. - - :rtype: flytekit.common.tasks.sdk_runnable.SdkRunnableTask - """ - - def wrapper(fn): - return (cls or _sdk_runnable_tasks.SdkRunnableTask)( - task_function=fn, - task_type=_common_constants.SdkTaskType.PYTHON_TASK, - discovery_version=cache_version, - retries=retries, - interruptible=interruptible, - deprecated=deprecated, - storage_request=storage_request, - cpu_request=cpu_request, - gpu_request=gpu_request, - memory_request=memory_request, - storage_limit=storage_limit, - cpu_limit=cpu_limit, - gpu_limit=gpu_limit, - memory_limit=memory_limit, - discoverable=cache, - timeout=timeout or _datetime.timedelta(seconds=0), - environment=environment, - cache_serializable=cache_serialize, - custom={}, - ) - - if _task_function: - return wrapper(_task_function) - else: - return wrapper - - -def dynamic_task( - _task_function=None, - cache_version="", - retries=0, - interruptible=None, - deprecated="", - storage_request=None, - cpu_request=None, - gpu_request=None, - memory_request=None, - storage_limit=None, - cpu_limit=None, - gpu_limit=None, - memory_limit=None, - cache=False, - timeout=None, - allowed_failure_ratio=None, - max_concurrency=None, - environment=None, - cache_serialize=False, - cls=None, -): - """ - Decorator to create a custom dynamic task definition. Dynamic tasks should be used to split up work into - an arbitrary number of parallel sub-tasks, or workflows. - - .. code-block:: python - - @outputs(out=Types.Integer) - @python_task - def my_sub_task(wf_params, out): - out.set(randint()) - - @outputs(out=[Types.Integer]) - @dynamic_task - def my_task(wf_params, out): - out_list = [] - for i in xrange(100): - out_list.append(my_sub_task().outputs.out) - out.set(out_list) - - .. note:: - - All outputs of a batch task must be a list. This is because the individual outputs of sub-tasks should be - appended into a list. There cannot be aggregation of outputs done in this task. To accomplish aggregation, - it is recommended that a python_task take the outputs of this task as input and do the necessary work. - If a sub-task does not contribute an output, it must be yielded from the task with the `yield` keyword or - returned from the task in a list. If this isn't done, the sub-task will not be executed. - - :param _task_function: this is the decorated method and shouldn't be declared explicitly. The function must - take a first argument, and then named arguments matching those defined in @inputs and @outputs. No keyword - arguments are allowed. - :param Text cache_version: [optional] string representing logical version for discovery. This field should be - updated whenever the underlying algorithm changes. - - .. note:: - - This argument is required to be a non-empty string if `cache` is True. - - :param int retries: [optional] integer determining number of times task can be retried on - :py:exc:`flytekit.sdk.exceptions.RecoverableException` or transient platform failures. Defaults - to 0. - - .. note:: - - If retries > 0, the task must be able to recover from any remote state created within the user code. It is - strongly recommended that tasks are written to be idempotent. - - :param bool interruptible: [optional] boolean describing if the task is interruptible. - - :param Text deprecated: [optional] string that should be provided if this task is deprecated. The string - will be logged as a warning so it should contain information regarding how to update to a newer task. - :param Text storage_request: [optional] Kubernetes resource string for lower-bound of disk storage space - for the task to run. Default is set by platform-level configuration. - - .. note:: - - This is currently not supported by the platform. - - :param Text cpu_request: [optional] Kubernetes resource string for lower-bound of cores for the task to execute. - This can be set to a fractional portion of a CPU. Default is set by platform-level configuration. - TODO: Add links to resource string documentation for Kubernetes - :param Text gpu_request: [optional] Kubernetes resource string for lower-bound of desired GPUs. - Default is set by platform-level configuration. - TODO: Add links to resource string documentation for Kubernetes - :param Text memory_request: [optional] Kubernetes resource string for lower-bound of physical memory - necessary for the task to execute. Default is set by platform-level configuration. - TODO: Add links to resource string documentation for Kubernetes - :param Text storage_limit: [optional] Kubernetes resource string for upper-bound of disk storage space - for the task to run. This amount is not guaranteed! If not specified, it is set equal to storage_request. - - .. note:: - - This is currently not supported by the platform. - - :param Text cpu_limit: [optional] Kubernetes resource string for upper-bound of cores for the task to execute. - This can be set to a fractional portion of a CPU. This amount is not guaranteed! If not specified, - it is set equal to cpu_request. - :param Text gpu_limit: [optional] Kubernetes resource string for upper-bound of desired GPUs. This amount is not - guaranteed! If not specified, it is set equal to gpu_request. - :param Text memory_limit: [optional] Kubernetes resource string for upper-bound of physical memory - necessary for the task to execute. This amount is not guaranteed! If not specified, it is set equal to - memory_request. - :param bool cache: [optional] boolean describing if the outputs of this task should be cached and - re-usable. - :param datetime.timedelta timeout: [optional] describes how long the task should be allowed to - run at max before triggering a retry (if retries are enabled). By default, tasks are allowed to run - indefinitely. If a null timedelta is passed (i.e. timedelta(seconds=0)), the task will not timeout. - :param float allowed_failure_ratio: [optional] float value describing the ratio of sub-tasks that may fail before - the master batch task considers itself a failure. By default, the value is 0 so if any task fails, the master - batch task will be marked a failure. If specified, the value must be between 0 and 1 inclusive. In the event a - non-zero value is specified, downstream tasks must be able to accept None values as outputs from individual - sub-tasks because the output values will be set to None for any sub-task that fails. - :param int max_concurrency: [optional] integer value describing the maximum number of tasks to run concurrently. - This is a stand-in pending better concurrency controls for special use-cases. The existence of this parameter - is not guaranteed between versions and therefore it is NOT recommended that it be used. - :param dict[Text,Text] environment: [optional] environment variables to set when executing this task. - :param bool cache_serialize: [optional] boolean describing if instances of this cachable task should be executed - in serial. This means only a single instances executes and other concurrent executions wait for it to complete - and reuse the cached outputs. - :param cls: This can be used to override the task implementation with a user-defined extension. The class - provided must be a subclass of flytekit.common.tasks.sdk_runnable.SdkRunnableTask. Generally, it should be a - subclass of flytekit.common.tasks.sdk_dynamic.SdkDynamicTask. A user can use this parameter to inject bespoke - logic into the base Flyte programming model. - :rtype: flytekit.common.tasks.sdk_runnable.SdkDynamicTask - """ - - def wrapper(fn): - return (cls or _sdk_dynamic.SdkDynamicTask)( - task_function=fn, - task_type=_common_constants.SdkTaskType.DYNAMIC_TASK, - discovery_version=cache_version, - retries=retries, - interruptible=interruptible, - deprecated=deprecated, - storage_request=storage_request, - cpu_request=cpu_request, - gpu_request=gpu_request, - memory_request=memory_request, - storage_limit=storage_limit, - cpu_limit=cpu_limit, - gpu_limit=gpu_limit, - memory_limit=memory_limit, - discoverable=cache, - timeout=timeout or _datetime.timedelta(seconds=0), - allowed_failure_ratio=allowed_failure_ratio, - max_concurrency=max_concurrency, - environment=environment or {}, - cache_serializable=cache_serialize, - custom={}, - ) - - if _task_function: - return wrapper(_task_function) - else: - return wrapper - - -def spark_task( - _task_function=None, - cache_version="", - retries=0, - interruptible=None, - deprecated="", - cache=False, - timeout=None, - spark_conf=None, - hadoop_conf=None, - environment=None, - cache_serialize=False, - cls=None, -): - """ - Decorator to create a spark task. This task will connect to a Spark cluster, configure the environment, - and then execute the code within the _task_function as the Spark driver program. - - .. code-block:: python - - @inputs(a=Types.Integer) - @spark_task( - spark_conf={ - 'spark.executor.cores': '7', - 'spark.executor.instances': '31', - 'spark.executor.memory': '32G' - } - ) - def sparky(wf_params, spark_context, a): - pass - - :param _task_function: this is the decorated method and shouldn't be declared explicitly. The function must - take a first argument, and then named arguments matching those defined in @inputs and @outputs. No keyword - arguments are allowed for wrapped task functions. - :param Text cache_version: [optional] string representing logical version for discovery. This field should be - updated whenever the underlying algorithm changes. - - .. note:: - - This argument is required to be a non-empty string if `cache` is True. - - :param int retries: [optional] integer determining number of times task can be retried on - :py:exc:`flytekit.sdk.exceptions.RecoverableException` or transient platform failures. Defaults - to 0. - - .. note:: - - If retries > 0, the task must be able to recover from any remote state created within the user code. It is - strongly recommended that tasks are written to be idempotent. - - :param Text deprecated: [optional] string that should be provided if this task is deprecated. The string - will be logged as a warning so it should contain information regarding how to update to a newer task. - :param bool cache: [optional] boolean describing if the outputs of this task should be cached and - re-usable. - :param datetime.timedelta timeout: [optional] describes how long the task should be allowed to - run at max before triggering a retry (if retries are enabled). By default, tasks are allowed to run - indefinitely. If a null timedelta is passed (i.e. timedelta(seconds=0)), the task will not timeout. - :param dict[Text,Text] spark_conf: A definition of key-value pairs for spark config for the job. - :param dict[Text,Text] hadoop_conf: A definition of key-value pairs for hadoop config for the job. - :param dict[Text,Text] environment: [optional] environment variables to set when executing this task. - :param bool cache_serialize: [optional] boolean describing if instances of this cachable task should be executed - in serial. This means only a single instances executes and other concurrent executions wait for it to complete - and reuse the cached outputs. - :param cls: This can be used to override the task implementation with a user-defined extension. The class - provided must be a subclass of flytekit.common.tasks.sdk_runnable.SdkRunnableTask. Generally, it should be a - subclass of flytekit.common.tasks.spark_task.SdkSparkTask. A user can use this parameter to inject bespoke - logic into the base Flyte programming model. - :rtype: flytekit.common.tasks.sdk_runnable.SdkRunnableTask - """ - - def wrapper(fn): - return (cls or _sdk_spark_tasks.SdkSparkTask)( - task_function=fn, - task_type=_common_constants.SdkTaskType.SPARK_TASK, - discovery_version=cache_version, - retries=retries, - interruptible=interruptible, - spark_type=_spark_type.PYTHON, - deprecated=deprecated, - discoverable=cache, - timeout=timeout or _datetime.timedelta(seconds=0), - spark_conf=spark_conf or {}, - hadoop_conf=hadoop_conf or {}, - environment=environment or {}, - cache_serializable=cache_serialize, - ) - - if _task_function: - return wrapper(_task_function) - else: - return wrapper - - -def generic_spark_task( - spark_type, - main_class, - main_application_file, - cache_version="", - retries=0, - interruptible=None, - inputs=None, - deprecated="", - cache=False, - timeout=None, - spark_conf=None, - hadoop_conf=None, - environment=None, - cache_serialize=False, -): - """ - Create a generic spark task. This task will connect to a Spark cluster, configure the environment, - and then execute the mainClass code as the Spark driver program. - - """ - - return _sdk_generic_spark_task.SdkGenericSparkTask( - task_type=_common_constants.SdkTaskType.SPARK_TASK, - discovery_version=cache_version, - retries=retries, - interruptible=interruptible, - deprecated=deprecated, - discoverable=cache, - timeout=timeout or _datetime.timedelta(seconds=0), - spark_type=spark_type, - task_inputs=inputs, - main_class=main_class or "", - main_application_file=main_application_file or "", - spark_conf=spark_conf or {}, - hadoop_conf=hadoop_conf or {}, - environment=environment or {}, - cache_serializable=cache_serialize, - ) - - -def qubole_spark_task(*args, **kwargs): - """ - :rtype: flytekit.common.tasks.sdk_runnable.SdkRunnableTask - """ - raise NotImplementedError("Qubole Spark Task is currently not supported in Flyte.") - - -def hive_task( - _task_function=None, - cache_version="", - retries=0, - interruptible=None, - deprecated="", - storage_request=None, - cpu_request=None, - gpu_request=None, - memory_request=None, - storage_limit=None, - cpu_limit=None, - gpu_limit=None, - memory_limit=None, - cache=False, - timeout=None, - environment=None, - cache_serialize=False, - cls=None, -): - """ - Decorator to create a hive task. This task should output a list of hive queries which are run on a hive cluster. - - This is a 2 step task: - - 1. Generator step runs in the user container and outputs a list of queries. - 2. The list of queries produced in step1 is then submitted to Hive Cluster. The queries are monitored by Flyte - Backend for completion. - Container properties(cpu, gpu, memory, etc) set on this task are only used in step1 above. - - .. code-block:: python - - @inputs(a=Types.Integer) - @hive_task( - cache_version='1', - ) - def test_hive(wf_params, a): - return [ - "SELECT * FROM users_table WHERE user_id=4", - "INSERT INTO users_table VALUES ("user", 5, 4)" - ] - - :param _task_function: this is the decorated method and shouldn't be declared explicitly. The function must - take a first argument, and then named arguments matching those defined in @inputs and @outputs. No keyword - arguments are allowed for wrapped task functions. - - :param Text cache_version: [optional] string representing logical version for discovery. This field should be - updated whenever the underlying algorithm changes. - - .. note:: - - This argument is required to be a non-empty string if `cache` is True. - - :param int retries: [optional] integer determining number of times task can be retried on - :py:exc:`flytekit.common.exceptions.RecoverableException` or transient platform failures. Defaults - to 0. - - .. note:: - - If retries > 0, the task must be able to recover from any remote state created within the user code. It is - strongly recommended that tasks are written to be idempotent. - - :param bool interruptible: [optional] boolean describing if task is interruptible. - :param Text deprecated: [optional] string that should be provided if this task is deprecated. The string - will be logged as a warning so it should contain information regarding how to update to a newer task. - :param Text storage_request: [optional] Kubernetes resource string for lower-bound of disk storage space - for the task to run. Default is set by platform-level configuration. - - .. note:: - - This is currently not supported by the platform. - - :param Text cpu_request: [optional] Kubernetes resource string for lower-bound of cores for the task to execute. - This can be set to a fractional portion of a CPU. Default is set by platform-level configuration. - TODO: Add links to resource string documentation for Kubernetes - :param Text gpu_request: [optional] Kubernetes resource string for lower-bound of desired GPUs. - Default is set by platform-level configuration. - TODO: Add links to resource string documentation for Kubernetes - :param Text memory_request: [optional] Kubernetes resource string for lower-bound of physical memory - necessary for the task to execute. Default is set by platform-level configuration. - TODO: Add links to resource string documentation for Kubernetes - :param Text storage_limit: [optional] Kubernetes resource string for upper-bound of disk storage space - for the task to run. This amount is not guaranteed! If not specified, it is set equal to storage_request. - - .. note:: - - This is currently not supported by the platform. - - :param Text cpu_limit: [optional] Kubernetes resource string for upper-bound of cores for the task to execute. - This can be set to a fractional portion of a CPU. This amount is not guaranteed! If not specified, - it is set equal to cpu_request. - :param Text gpu_limit: [optional] Kubernetes resource string for upper-bound of desired GPUs. This amount is not - guaranteed! If not specified, it is set equal to gpu_request. - :param Text memory_limit: [optional] Kubernetes resource string for upper-bound of physical memory - necessary for the task to execute. This amount is not guaranteed! If not specified, it is set equal to - memory_request. - :param bool cache: [optional] boolean describing if the outputs of this task should be cached and - re-usable. - :param datetime.timedelta timeout: [optional] describes how long the task should be allowed to - run at max before triggering a retry (if retries are enabled). By default, tasks are allowed to run - indefinitely. If a null timedelta is passed (i.e. timedelta(seconds=0)), the task will not timeout. - :param dict[Text,Text] environment: Environment variables to set for the execution of the query-generating - container. - :param bool cache_serialize: [optional] boolean describing if instances of this cachable task should be executed - in serial. This means only a single instances executes and other concurrent executions wait for it to complete - and reuse the cached outputs. - :param cls: This can be used to override the task implementation with a user-defined extension. The class - provided should be a subclass of flytekit.common.tasks.sdk_runnable.SdkRunnableTask. Generally, it should be - a subclass of flytekit.common.tasks.hive_task.SdkHiveTask. A user can use this to inject bespoke logic into - the base Flyte programming model. - - :rtype: flytekit.common.tasks.sdk_runnable.SdkHiveTask - """ - - def wrapper(fn): - - return (cls or _sdk_hive_tasks.SdkHiveTask)( - task_function=fn, - task_type=_common_constants.SdkTaskType.BATCH_HIVE_TASK, - discovery_version=cache_version, - retries=retries, - interruptible=interruptible, - deprecated=deprecated, - storage_request=storage_request, - cpu_request=cpu_request, - gpu_request=gpu_request, - memory_request=memory_request, - storage_limit=storage_limit, - cpu_limit=cpu_limit, - gpu_limit=gpu_limit, - memory_limit=memory_limit, - discoverable=cache, - timeout=timeout or _datetime.timedelta(seconds=0), - cluster_label="", - tags=[], - environment=environment or {}, - cache_serializable=cache_serialize, - ) - - if _task_function: - return wrapper(_task_function) - else: - return wrapper - - -def qubole_hive_task( - _task_function=None, - cache_version="", - retries=0, - interruptible=None, - deprecated="", - storage_request=None, - cpu_request=None, - gpu_request=None, - memory_request=None, - storage_limit=None, - cpu_limit=None, - gpu_limit=None, - memory_limit=None, - cache=False, - timeout=None, - cluster_label=None, - tags=None, - environment=None, - cache_serialize=False, - cls=None, -): - """ - Decorator to create a qubole hive task. This is hive task runs on a qubole cluster, and therefore allows users to - pass cluster labels and qubole query tags. Similar to hive task, this task should output a list of hive queries - that are run on a hive cluster. - - Similar to a hive task, this is also a 2 step task where step2 is run on a qubole hive cluster. Therefore, users can - specify qubole cluster_label and query tags on this task. - - .. code-block:: python - - @inputs(a=Types.Integer) - @qubole_hive_task( - cache_version='1', - cluster_label='cluster_label', - tags=['tag1'], - ) - def test_hive(wf_params, a): - return [ - "SELECT * FROM users_table WHERE user_id=4", - "INSERT INTO users_table VALUES ("user", 5, 4)" - ] - - :param _task_function: this is the decorated method and shouldn't be declared explicitly. The function must - take a first argument, and then named arguments matching those defined in @inputs and @outputs. No keyword - arguments are allowed for wrapped task functions. - - :param Text cache_version: [optional] string representing logical version for discovery. This field should be - updated whenever the underlying algorithm changes. - - .. note:: - - This argument is required to be a non-empty string if `cache` is True. - - :param int retries: [optional] integer determining number of times task can be retried on - :py:exc:`flytekit.common.exceptions.RecoverableException` or transient platform failures. Defaults - to 0. - - .. note:: - - If retries > 0, the task must be able to recover from any remote state created within the user code. It is - strongly recommended that tasks are written to be idempotent. - - :param bool interruptible: [optional] boolean describing if task is interruptible. - :param Text deprecated: [optional] string that should be provided if this task is deprecated. The string - will be logged as a warning so it should contain information regarding how to update to a newer task. - :param Text storage_request: [optional] Kubernetes resource string for lower-bound of disk storage space - for the task to run. Default is set by platform-level configuration. - - .. note:: - - This is currently not supported by the platform. - - :param Text cpu_request: [optional] Kubernetes resource string for lower-bound of cores for the task to execute. - This can be set to a fractional portion of a CPU. Default is set by platform-level configuration. - TODO: Add links to resource string documentation for Kubernetes - :param Text gpu_request: [optional] Kubernetes resource string for lower-bound of desired GPUs. - Default is set by platform-level configuration. - TODO: Add links to resource string documentation for Kubernetes - :param Text memory_request: [optional] Kubernetes resource string for lower-bound of physical memory - necessary for the task to execute. Default is set by platform-level configuration. - TODO: Add links to resource string documentation for Kubernetes - :param Text storage_limit: [optional] Kubernetes resource string for upper-bound of disk storage space - for the task to run. This amount is not guaranteed! If not specified, it is set equal to storage_request. - - .. note:: - - This is currently not supported by the platform. - - :param Text cpu_limit: [optional] Kubernetes resource string for upper-bound of cores for the task to execute. - This can be set to a fractional portion of a CPU. This amount is not guaranteed! If not specified, - it is set equal to cpu_request. - :param Text gpu_limit: [optional] Kubernetes resource string for upper-bound of desired GPUs. This amount is not - guaranteed! If not specified, it is set equal to gpu_request. - :param Text memory_limit: [optional] Kubernetes resource string for upper-bound of physical memory - necessary for the task to execute. This amount is not guaranteed! If not specified, it is set equal to - memory_request. - :param bool cache: [optional] boolean describing if the outputs of this task should be cached and - re-usable. - :param datetime.timedelta timeout: [optional] describes how long the task should be allowed to - run at max before triggering a retry (if retries are enabled). By default, tasks are allowed to run - indefinitely. If a null timedelta is passed (i.e. timedelta(seconds=0)), the task will not timeout. - :param cluster_label: The qubole cluster label where the query is to be executed - :param list[Text] tags: User defined tags(key-value pairs) defined by the user for the queries. These tags are - passed to Qubole. - :param dict[Text,Text] environment: Environment variables to set for the execution of the query-generating - container. - :param bool cache_serialize: [optional] boolean describing if instances of this cachable task should be executed - in serial. This means only a single instances executes and other concurrent executions wait for it to complete - and reuse the cached outputs. - :param cls: This can be used to override the task implementation with a user-defined extension. The class - provided should be a subclass of flytekit.common.tasks.sdk_runnable.SdkRunnableTask. Generally, it should be - a subclass of flytekit.common.tasks.hive_task.SdkHiveTask. A user can use this to inject bespoke logic into - the base Flyte programming model. - - :rtype: flytekit.common.tasks.sdk_runnable.SdkHiveTask - """ - - def wrapper(fn): - - return (cls or _sdk_hive_tasks.SdkHiveTask)( - task_function=fn, - task_type=_common_constants.SdkTaskType.BATCH_HIVE_TASK, - discovery_version=cache_version, - retries=retries, - interruptible=interruptible, - deprecated=deprecated, - storage_request=storage_request, - cpu_request=cpu_request, - gpu_request=gpu_request, - memory_request=memory_request, - storage_limit=storage_limit, - cpu_limit=cpu_limit, - gpu_limit=gpu_limit, - memory_limit=memory_limit, - discoverable=cache, - timeout=timeout or _datetime.timedelta(seconds=0), - cluster_label=cluster_label or "", - tags=tags or [], - environment=environment or {}, - cache_serializable=cache_serialize, - ) - - # This is syntactic-sugar, so that when calling this decorator without args, you can either - # do it with () or without any () - if _task_function: - return wrapper(_task_function) - else: - return wrapper - - -def sidecar_task( - _task_function=None, - cache_version="", - retries=0, - interruptible=None, - deprecated="", - storage_request=None, - cpu_request=None, - gpu_request=None, - memory_request=None, - storage_limit=None, - cpu_limit=None, - gpu_limit=None, - memory_limit=None, - cache=False, - timeout=None, - environment=None, - cache_serialize=False, - pod_spec=None, - primary_container_name=None, - annotations=None, - labels=None, - cls=None, -): - """ - Decorator to create a Sidecar Task definition. This task will execute the primary task alongside the specified - kubernetes PodSpec. Custom primary task container attributes can be defined in the PodSpec by defining a container - whose name matches the primary_container_name. These container attributes will be applied to the container brought - up to execute the primary task definition. - - .. code-block:: python - - def generate_pod_spec_for_task(): - pod_spec = generated_pb2.PodSpec() - secondary_container = generated_pb2.Container( - name="secondary", - image="alpine", - ) - secondary_container.command.extend(["/bin/sh"]) - secondary_container.args.extend(["-c", "echo hi sidecar world > /data/message.txt"]) - shared_volume_mount = generated_pb2.VolumeMount( - name="shared-data", - mountPath="/data", - ) - secondary_container.volumeMounts.extend([shared_volume_mount]) - - primary_container = generated_pb2.Container(name="primary") - primary_container.volumeMounts.extend([shared_volume_mount]) - - pod_spec.volumes.extend([generated_pb2.Volume( - name="shared-data", - volumeSource=generated_pb2.VolumeSource( - emptyDir=generated_pb2.EmptyDirVolumeSource( - medium="Memory", - ) - ) - )]) - pod_spec.containers.extend([primary_container, secondary_container]) - return pod_spec - - @sidecar_task( - pod_spec=generate_pod_spec_for_task(), - primary_container_name="primary", - annotations={"key": "value"}, - labels={"key": "value"}, - ) - def a_sidecar_task(wfparams): - while not os.path.isfile('/data/message.txt'): - time.sleep(5) - - - :param _task_function: this is the decorated method and shouldn't be declared explicitly. The function must - take a first argument, and then named arguments matching those defined in @inputs and @outputs. No keyword - arguments are allowed for wrapped task functions. - - :param Text cache_version: [optional] string representing logical version for discovery. This field should be - updated whenever the underlying algorithm changes. - .. note:: - This argument is required to be a non-empty string if `cache` is True. - - :param int retries: [optional] integer determining number of times task can be retried on - :py:ex:`flytekit.sdk.exceptions.RecoverableException` or transient platform failures. Defaults - to 0. - - .. note:: - - If retries > 0, the task must be able to recover from any remote state created within the user code. It is - strongly recommended that tasks are written to be idempotent. - - :param bool interruptible: Specify whether task is interruptible - - :param Text deprecated: [optional] string that should be provided if this task is deprecated. The string - will be logged as a warning so it should contain information regarding how to update to a newer task. - - :param Text storage_request: [optional] Kubernetes resource string for lower-bound of disk storage space - for the task to run. Default is set by platform-level configuration. - - TODO: !!! - .. note:: - - This is currently not supported by the platform. - - :param Text cpu_request: [optional] Kubernetes resource string for lower-bound of cores for the task to execute. - This can be set to a fractional portion of a CPU. Default is set by platform-level configuration. - - TODO: Add links to resource string documentation for Kubernetes - - :param Text gpu_request: [optional] Kubernetes resource string for lower-bound of desired GPUs. - Default is set by platform-level configuration. - - TODO: Add links to resource string documentation for Kubernetes - - :param Text memory_request: [optional] Kubernetes resource string for lower-bound of physical memory - necessary for the task to execute. Default is set by platform-level configuration. - - TODO: Add links to resource string documentation for Kubernetes - - :param Text storage_limit: [optional] Kubernetes resource string for upper-bound of disk storage space - for the task to run. This amount is not guaranteed! If not specified, it is set equal to storage_request. - - TODO: !!! - .. note:: - - This is currently not supported by the platform. - - :param Text cpu_limit: [optional] Kubernetes resource string for upper-bound of cores for the task to execute. - This can be set to a fractional portion of a CPU. This amount is not guaranteed! If not specified, - it is set equal to cpu_request. - - :param Text gpu_limit: [optional] Kubernetes resource string for upper-bound of desired GPUs. This amount is not - guaranteed! If not specified, it is set equal to gpu_request. - - :param Text memory_limit: [optional] Kubernetes resource string for upper-bound of physical memory - necessary for the task to execute. This amount is not guaranteed! If not specified, it is set equal to - memory_request. - - :param bool cache: [optional] boolean describing if the outputs of this task should be discoverable and - re-usable. - - :param datetime.timedelta timeout: [optional] describes how long the task should be allowed to - run at max before triggering a retry (if retries are enabled). By default, tasks are allowed to run - indefinitely. If a null timedelta is passed (i.e. timedelta(seconds=0)), the task will not timeout. - - :param dict[Text,Text] environment: [optional] environment variables to set when executing this task. - - :param bool cache_serialize: [optional] boolean describing if instances of this cachable task should be executed - in serial. This means only a single instances executes and other concurrent executions wait for it to complete - and reuse the cached outputs. - - :param k8s.io.api.core.v1.generated_pb2.PodSpec pod_spec: [optional] PodSpec to bring up alongside task execution. - - :param Text primary_container_name: primary container to monitor for the duration of the task. - - :param dict[Text, Text] annotations: [optional] kubernetes annotations - - :param dict[Text, Text] labels: [optional] kubernetes labels - - :param cls: This can be used to override the task implementation with a user-defined extension. The class - provided must be a subclass of flytekit.common.tasks.sdk_runnable.SdkRunnableTask. A user can use this to - inject bespoke logic into the base Flyte programming model. - - :rtype: flytekit.common.tasks.sdk_runnable.SdkRunnableTask - - """ - - def wrapper(fn): - - return (cls or _sdk_sidecar_tasks.SdkSidecarTask)( - task_function=fn, - task_type=_common_constants.SdkTaskType.SIDECAR_TASK, - discovery_version=cache_version, - retries=retries, - interruptible=interruptible, - deprecated=deprecated, - storage_request=storage_request, - cpu_request=cpu_request, - gpu_request=gpu_request, - memory_request=memory_request, - storage_limit=storage_limit, - cpu_limit=cpu_limit, - gpu_limit=gpu_limit, - memory_limit=memory_limit, - discoverable=cache, - timeout=timeout or _datetime.timedelta(seconds=0), - environment=environment, - cache_serializable=cache_serialize, - pod_spec=pod_spec, - primary_container_name=primary_container_name, - annotations=annotations, - labels=labels, - ) - - if _task_function: - return wrapper(_task_function) - else: - return wrapper - - -def dynamic_sidecar_task( - _task_function=None, - cache_version="", - retries=0, - interruptible=None, - deprecated="", - storage_request=None, - cpu_request=None, - gpu_request=None, - memory_request=None, - storage_limit=None, - cpu_limit=None, - gpu_limit=None, - memory_limit=None, - cache=False, - timeout=None, - allowed_failure_ratio=None, - max_concurrency=None, - environment=None, - cache_serialize=False, - pod_spec=None, - primary_container_name=None, - annotations=None, - labels=None, - cls=None, -): - """ - Decorator to create a custom dynamic sidecar task definition. Dynamic - tasks should be used to split up work into an arbitrary number of parallel - sub-tasks, or workflows. This task will execute the primary task alongside - the specified kubernetes PodSpec. Custom primary task container attributes - can be defined in the PodSpec by defining a container whose name matches - the primary_container_name. These container attributes will be applied to - the container brought up to execute the primary task definition. - .. code-block:: python - def generate_pod_spec_for_task(): - pod_spec = generated_pb2.PodSpec() - secondary_container = generated_pb2.Container( - name="secondary", - image="alpine", - ) - secondary_container.command.extend(["/bin/sh"]) - secondary_container.args.extend(["-c", "echo hi sidecar world > /data/message.txt"]) - shared_volume_mount = generated_pb2.VolumeMount( - name="shared-data", - mountPath="/data", - ) - secondary_container.volumeMounts.extend([shared_volume_mount]) - primary_container = generated_pb2.Container(name="primary") - primary_container.volumeMounts.extend([shared_volume_mount]) - pod_spec.volumes.extend([generated_pb2.Volume( - name="shared-data", - volumeSource=generated_pb2.VolumeSource( - emptyDir=generated_pb2.EmptyDirVolumeSource( - medium="Memory", - ) - ) - )]) - pod_spec.containers.extend([primary_container, secondary_container]) - return pod_spec - @outputs(out=Types.Integer) - @python_task - def my_sub_task(wf_params, out): - out.set(randint()) - @outputs(out=[Types.Integer]) - @dynamic_sidecar_task( - pod_spec=generate_pod_spec_for_task(), - primary_container_name="primary", - annotations={"a": "a"}, - labels={"b": "b"}, - ) - def my_task(wf_params, out): - out_list = [] - for i in xrange(100): - out_list.append(my_sub_task().outputs.out) - out.set(out_list) - .. note:: - All outputs of a batch task must be a list. This is because the individual outputs of sub-tasks should be - appended into a list. There cannot be aggregation of outputs done in this task. To accomplish aggregation, - it is recommended that a python_task take the outputs of this task as input and do the necessary work. - If a sub-task does not contribute an output, it must be yielded from the task with the `yield` keyword or - returned from the task in a list. If this isn't done, the sub-task will not be executed. - :param _task_function: this is the decorated method and shouldn't be declared explicitly. The function must - take a first argument, and then named arguments matching those defined in @inputs and @outputs. No keyword - arguments are allowed. - :param Text cache_version: [optional] string representing logical version for discovery. This field should be - updated whenever the underlying algorithm changes. - .. note:: - This argument is required to be a non-empty string if `cache` is True. - :param int retries: [optional] integer determining number of times task can be retried on - :py:exc:`flytekit.sdk.exceptions.RecoverableException` or transient platform failures. Defaults - to 0. - .. note:: - If retries > 0, the task must be able to recover from any remote state created within the user code. It is - strongly recommended that tasks are written to be idempotent. - :param bool interruptible: [optional] boolean describing if the task is interruptible. - :param Text deprecated: [optional] string that should be provided if this task is deprecated. The string - will be logged as a warning so it should contain information regarding how to update to a newer task. - :param Text storage_request: [optional] Kubernetes resource string for lower-bound of disk storage space - for the task to run. Default is set by platform-level configuration. - .. note:: - This is currently not supported by the platform. - :param Text cpu_request: [optional] Kubernetes resource string for lower-bound of cores for the task to execute. - This can be set to a fractional portion of a CPU. Default is set by platform-level configuration. - TODO: Add links to resource string documentation for Kubernetes - :param Text gpu_request: [optional] Kubernetes resource string for lower-bound of desired GPUs. - Default is set by platform-level configuration. - TODO: Add links to resource string documentation for Kubernetes - :param Text memory_request: [optional] Kubernetes resource string for lower-bound of physical memory - necessary for the task to execute. Default is set by platform-level configuration. - TODO: Add links to resource string documentation for Kubernetes - :param Text storage_limit: [optional] Kubernetes resource string for upper-bound of disk storage space - for the task to run. This amount is not guaranteed! If not specified, it is set equal to storage_request. - .. note:: - This is currently not supported by the platform. - :param Text cpu_limit: [optional] Kubernetes resource string for upper-bound of cores for the task to execute. - This can be set to a fractional portion of a CPU. This amount is not guaranteed! If not specified, - it is set equal to cpu_request. - :param Text gpu_limit: [optional] Kubernetes resource string for upper-bound of desired GPUs. This amount is not - guaranteed! If not specified, it is set equal to gpu_request. - :param Text memory_limit: [optional] Kubernetes resource string for upper-bound of physical memory - necessary for the task to execute. This amount is not guaranteed! If not specified, it is set equal to - memory_request. - :param bool cache: [optional] boolean describing if the outputs of this task should be cached and - re-usable. - :param datetime.timedelta timeout: [optional] describes how long the task should be allowed to - run at max before triggering a retry (if retries are enabled). By default, tasks are allowed to run - indefinitely. If a null timedelta is passed (i.e. timedelta(seconds=0)), the task will not timeout. - :param float allowed_failure_ratio: [optional] float value describing the ratio of sub-tasks that may fail before - the master batch task considers itself a failure. By default, the value is 0 so if any task fails, the master - batch task will be marked a failure. If specified, the value must be between 0 and 1 inclusive. In the event a - non-zero value is specified, downstream tasks must be able to accept None values as outputs from individual - sub-tasks because the output values will be set to None for any sub-task that fails. - :param int max_concurrency: [optional] integer value describing the maximum number of tasks to run concurrently. - This is a stand-in pending better concurrency controls for special use-cases. The existence of this parameter - is not guaranteed between versions and therefore it is NOT recommended that it be used. - :param dict[Text,Text] environment: [optional] environment variables to set when executing this task. - :param bool cache_serialize: [optional] boolean describing if instances of this cachable task should be executed - in serial. This means only a single instances executes and other concurrent executions wait for it to complete - and reuse the cached outputs. - :param k8s.io.api.core.v1.generated_pb2.PodSpec pod_spec: PodSpec to bring up alongside task execution. - :param Text primary_container_name: primary container to monitor for the duration of the task. - :param dict[Text, Text] annotations: [optional] kubernetes annotations - :param dict[Text, Text] labels: [optional] kubernetes labels - :param cls: This can be used to override the task implementation with a user-defined extension. The class - provided must be a subclass of flytekit.common.tasks.sdk_runnable.SdkRunnableTask. Generally, it should be a - subclass of flytekit.common.tasks.sidecar_Task.SdkDynamicSidecarTask. A user can use this parameter to inject bespoke - logic into the base Flyte programming model. - :rtype: flytekit.common.tasks.sidecar_Task.SdkDynamicSidecarTask - """ - - def wrapper(fn): - return (cls or _sdk_sidecar_tasks.SdkDynamicSidecarTask)( - task_function=fn, - task_type=_common_constants.SdkTaskType.SIDECAR_TASK, - discovery_version=cache_version, - retries=retries, - interruptible=interruptible, - deprecated=deprecated, - storage_request=storage_request, - cpu_request=cpu_request, - gpu_request=gpu_request, - memory_request=memory_request, - storage_limit=storage_limit, - cpu_limit=cpu_limit, - gpu_limit=gpu_limit, - memory_limit=memory_limit, - discoverable=cache, - timeout=timeout or _datetime.timedelta(seconds=0), - allowed_failure_ratio=allowed_failure_ratio, - max_concurrency=max_concurrency, - environment=environment, - cache_serializable=cache_serialize, - pod_spec=pod_spec, - primary_container_name=primary_container_name, - annotations=annotations, - labels=labels, - ) - - if _task_function: - return wrapper(_task_function) - else: - return wrapper diff --git a/flytekit/sdk/test_utils.py b/flytekit/sdk/test_utils.py deleted file mode 100644 index 766774bf2c..0000000000 --- a/flytekit/sdk/test_utils.py +++ /dev/null @@ -1,92 +0,0 @@ -from wrapt import decorator as _decorator - -from flytekit.common import utils as _utils -from flytekit.interfaces.data import data_proxy as _data_proxy - - -class LocalTestFileSystem(object): - """ - Context manager for creating a temporary test file system locally for the purpose of unit testing and grabbing - remote objects. This scratch space will be automatically cleaned up as long as sys.exit() is not called from - within the context. This context need only be used in user scripts and tests--all task executions are guaranteed - to have the necessary managed disk context available. - - .. note:: - - This is especially useful when dealing with remote blob-like objects (Blob, CSV, MultiPartBlob, - MultiPartCSV, Schema) as they require backing on disk. Using this context manager creates that disk context - to support the downloads. - - .. note:: - - Blob-like objects can be downloaded to user-specified locations. See documentation for - flytekit.sdk.types.Types for more information. - - .. note:: - - When this context is entered, it overrides any contexts already entered. All blobs will be written to the most - recent entered context. Upon exiting the context, all data associated will be deleted. It is recommended to - only use one LocalTestFileSystem() per test to avoid confusion. - - .. code-block:: python - - with LocalTestFileSystem(): - wf_handle = SdkWorkflowExecution.fetch('project', 'domain', 'name') - with wf_handle.node_executions['my_node'].outputs.blob as reader: - assert reader.read() == "hello!" - """ - - def __init__(self): - self._exit_stack = _utils.ExitStack() - - def __enter__(self): - """ - :rtype: flytekit.common.utils.AutoDeletingTempDir - """ - self._exit_stack.__enter__() - temp_dir = self._exit_stack.enter_context(_utils.AutoDeletingTempDir("local_test_filesystem")) - self._exit_stack.enter_context(_data_proxy.LocalDataContext(temp_dir.name)) - self._exit_stack.enter_context(_data_proxy.LocalWorkingDirectoryContext(temp_dir)) - return temp_dir - - def __exit__(self, exc_type, exc_val, exc_tb): - return self._exit_stack.__exit__(exc_type, exc_val, exc_tb) - - -@_decorator -def flyte_test(fn, _, args, kwargs): - """ - This is a decorator which can be used to annotate test functions. By using this decorator, the necessary local - scratch context will be prepared and then cleaned up upon completion. - - .. code-block:: python - - @inputs(input_blob=Types.Blob) - @outputs(response_blob=Types.Blob) - @python_task - def test_task(wf_params, input_blob, response_blob): - response = Types.Blob() - with response as writer: - with input_blob as reader: - txt = reader.read() - if txt == "Hi": - writer.write("Hello, world!") - elif txt == "Goodnight": - writer.write("Goodnight, moon.") - else: - writer.write("Does not compute".) - response_blob.set(response) - - @flyte_test - def some_test(): - blob = Types.Blob() - with blob as writer: - writer.write("Hi") - - result = test_task.unit_test(input_blob=blob) - - with result['response_blob'] as reader: - assert reader.read() == 'Hello, world!" - """ - with LocalTestFileSystem(): - return fn(*args, **kwargs) diff --git a/flytekit/sdk/types.py b/flytekit/sdk/types.py deleted file mode 100644 index 7b64d1db49..0000000000 --- a/flytekit/sdk/types.py +++ /dev/null @@ -1,504 +0,0 @@ -from flytekit.common.types import blobs as _blobs -from flytekit.common.types import containers as _containers -from flytekit.common.types import helpers as _helpers -from flytekit.common.types import primitives as _primitives -from flytekit.common.types import proto as _proto -from flytekit.common.types import schema as _schema - - -class Types(object): - Integer = _helpers.get_sdk_type_from_literal_type(_primitives.Integer.to_flyte_literal_type()) - """ - Use this to specify a simple integer type. - - When used with an SDK-decorated method, expect this behavior from the default type engine: - - As input: - 1) If set, A Python int will be received, if set. - 2) Otherwise, a None value will be received. - - As output: - 1) User code may pass an int or long value. - 2) Output can also be nulled with a None value. - - From command-line: - Specify an integer or integer string. - - .. code-block:: python - - @inputs(a=Types.Integer) - @outputs(b=Types.Integer) - @python_task - def double(wf_params, a, b): - b.set(a * 2) - """ - - Float = _helpers.get_sdk_type_from_literal_type(_primitives.Float.to_flyte_literal_type()) - """ - Use this to specify a simple floating point type. - - When used with an SDK-decorated method, expect this behavior from the default type engine: - - As input: - A Python float will be received, if set. Otherwise, a None value will be received. - - As output: - User code may pass a float value. It can also be nulled with a None value. - - From command-line: - Specify a float or floating-point string. - - .. code-block:: python - - @inputs(a=Types.Float) - @outputs(b=Types.Float) - @python_task - def square(wf_params, a, b): - b.set(a * a) - """ - - String = _helpers.get_sdk_type_from_literal_type(_primitives.String.to_flyte_literal_type()) - """ - Use this to specify a simple string type. - - When used with an SDK-decorated method, expect this behavior from the default type engine: - - As input: - A Python str (Python 2) or unicode (Python 3) will be received, if set. Otherwise, a None value will be - received. - - As output: - User code may pass a str value (Python 2) or a unicode value (Python 3). It can also be nulled with a None - value. - - From command-line: - Specify a string. - - .. code-block:: python - - @inputs(a=Types.String, b=Types.String) - @outputs(c=Types.String) - @python_task - def concat(wf_params, a, b): - c.set(a + b) - """ - - Boolean = _helpers.get_sdk_type_from_literal_type(_primitives.Boolean.to_flyte_literal_type()) - """ - Use this to specify a simple bool type. - - When used with an SDK-decorated method, expect this behavior from the default type engine: - - As input: - A Python bool will be received, if set. Otherwise, a None value will be received. - - As output: - User code may pass a bool value. It can also be nulled with a None value. - - From command-line: - Specify 0, 1, true, or false. - - .. code-block:: python - - @inputs(a=Types.Boolean) - @outputs(b=Types.Boolean) - @python_task - def invert(wf_params, a, b): - b.set(not a) - """ - - Datetime = _helpers.get_sdk_type_from_literal_type(_primitives.Datetime.to_flyte_literal_type()) - """ - Use this to specify a simple datetime type. - - When used with an SDK-decorated method, expect this behavior from the default type engine: - - As input: - A Python timezone-aware datetime.datetime will be received with a UTC time, if set. Otherwise, - a None value will be received. - - As output: - User code may pass a timezone-aware datetime.datetime value. It can also be nulled with a None value. - - From command-line: - Specify a timezone-aware, parsable datestring. i.e. 2019-01-01T00:00+00:00 - - .. note:: - - The engine requires that datetimes be timezone aware. By default, Python datetime.datetime is not timezone - aware. - - .. code-block:: python - - @inputs(a=Types.Datetime) - @outputs(b=Types.Datetime) - @python_task - def tomorrow(wf_params, a, b): - b.set(a + datetime.timedelta(days=1)) - """ - - Timedelta = _helpers.get_sdk_type_from_literal_type(_primitives.Timedelta.to_flyte_literal_type()) - """ - Use this to specify a simple timedelta type. - - When used with an SDK-decorated method, expect this behavior from the default type engine: - - As input: - A Python datetime.timedelta will be received, if set. Otherwise, a None value will be received. - - As output: - User code may pass a datetime.timedelta value. It can also be nulled with a None value. - - From command-line: - Specify a parsable duration string. i.e. 1h30m24s - - .. code-block:: python - - @inputs(a=Types.Timedelta) - @outputs(b=Types.Timedelta) - @python_task - def hundred_times_longer(wf_params, a, b): - b.set(a * 100) - """ - - Generic = _helpers.get_sdk_type_from_literal_type(_primitives.Generic.to_flyte_literal_type()) - """ - Use this to specify a simple JSON type. The Generic type offer a flexible (but loose) extension to flyte's typing - system by allowing custom types/objects to be passed through. It's strongly recommended for producers & consumers of - entities that produce or consume a Generic type to perform their own expectations checks on the integrity of the - object. - - When used with an SDK-decorated method, expect this behavior from the default type engine: - - As input: - 1) If set, a Python dict with JSON-ifiable primitives and nested lists or maps. - 2) Otherwise, a None value will be received. - - As output: - 1) User code may pass a Python dict with arbitrarily nested lists and dictionaries. JSON-ifiable - primitives may also be specified. - 2) Output can also be nulled with a None value. - - From command-line: - Specify a JSON string. - - .. code-block:: python - - @inputs(a=Types.Generic) - @outputs(b=Types.Generic) - @python_task - def operate(wf_params, a, b): - if a['operation'] == 'add': - a['value'] += a['operand'] # a['value'] is a number - elif a['operation'] == 'merge': - a['value'].update(a['some']['nested'][0]['field']) - b.set(a) - - # For better readability, it's strongly advised to leverage python's type aliasing. - MyTypeA = Types.Generic - MyTypeB = Types.Generic - - # This makes it clearer that it received a certain type and produces a different one. Other tasks that consume - # MyTypeB should do so in their input declaration. - @inputs(a=MyTypeA) - @outputs(b=MyTypeB) - @python_task - def operate(wf_params, a, b): - if a['operation'] == 'add': - a['value'] += a['operand'] # a['value'] is a number - elif a['operation'] == 'merge': - a['value'].update(a['some']['nested'][0]['field']) - b.set(a) - """ - - Blob = _blobs.Blob - """ - Use this to specify a Blob object which is essentially a managed file. - - When used with an SDK-decorated method, expect this behavior from the default type engine: - - As input: - 1) If set, a :py:class:`flytekit.common.types.impl.blobs.Blob` object will be received. - 2) If not set, a None value. - - As output: - 1) A user may specify a path string. - 2) A user may construct a :py:class:`flytekit.common.types.impl.blobs.Blob` object and pass it as output. - 3) Output can be nulled with a None value. - - From command-line: - Specify a path to the blob. This path must be accessible from the container when executing--either by - being downloaded from an accessible remote location like s3 or as a local file. - - .. code-block:: python - - @inputs(a=Types.Blob) - @outputs(b=Types.Blob) - @python_task - def copy(wf_params, a, b): - with a as reader: - txt = reader.read() - - out = Types.Blob() # Create at a random location specified in flytekit configuration - with out as writer: - writer.write(txt) - b.set(out) - """ - - CSV = _blobs.CSV - """ - Use this to specify a CSV blob object which is essentially a managed file in the CSV format. - - When used with an SDK-decorated method, expect this behavior from the default type engine: - - As input: - 1) If set, a :py:class:`flytekit.common.types.impl.blobs.CSV` object will be received. - 2) If not set, a None value. - - As output: - 1) A user may specify a path string. - 2) A user may construct a :py:class:`flytekit.common.types.impl.blobs.CSV` object and pass it as output. - 3) Output can be nulled with a None value. - - From command-line: - Specify a path to the CSV. This path must be accessible from the container when executing--either by - being downloaded from an accessible remote location like s3 or as a local file. - - .. code-block:: python - - @inputs(a=Types.CSV) - @outputs(b=Types.CSV) - @python_task - def copy(wf_params, a, b): - with a as reader: - txt = reader.read() - - out = Types.CSV() # Create at a random location specified in flytekit configuration - with out as writer: - writer.write(txt) - b.set(out) - """ - - MultiPartBlob = _blobs.MultiPartBlob - """ - Use this to specify a multi-part blob object which is essentially a chunked file in a non-recursive directory. - - When used with an SDK-decorated method, expect this behavior from the default type engine: - - As input: - 1) If set, a :py:class:`flytekit.common.types.impl.blobs.MultiPartBlob` object will be received. - 2) If not set, a None value. - - As output: - 1) A user may specify a path string. - 2) A user may construct a :py:class:`flytekit.common.types.impl.blobs.MultiPartBlob` object and pass it as - output. - 3) Output can be nulled with a None value. - - From command-line: - Specify a path to the multi-part blob. This path must be accessible from the container when - executing--either by being downloaded from an accessible remote location like s3 or as a local file. - - .. code-block:: python - - @inputs(a=Types.MultiPartBlob) - @outputs(b=Types.MultiPartBlob) - @python_task - def concat_then_split(wf_params, a, b): - txt = "" - with a as chunks: - for chunk in chunks: - txt += chunk.read() - - out = Types.MultiPartBlob() # Create at a random location specified in flytekit configuration - with out.create_part('000000') as writer: - writer.write("Chunk1") - with out.create_part('000001') as writer: - writer.write("Chunk2") - b.set(out) - """ - - MultiPartCSV = _blobs.MultiPartCSV - """ - See :py:attr:`flytekit.sdk.types.Types.MultiPartBlob`, but in CSV format - """ - - Schema = staticmethod(_schema.schema_instantiator) - """ - Use this to specify a Schema blob object which is essentially a chunked stream of Parquet dataframes. - - When used with an SDK-decorated method, expect this behavior from the default type engine: - - Cast behavior: - 1) A generic schema (specified as `Types.Schema()`) can receive input from any schema type regardless of - column definitions. - 2) A schema can receive as input any schema object as long as the upstream schema has a superset of the - column names defined and the types match for paired columns. Ordering does not matter. - - As input: - 1) If set, a :py:class:`flytekit.common.types.impl.schema.Schema` object will be received. - 2) If not set, a None value. - - As output: - 1) A user may specify a path string to a chunked dataframe non-recursive directory. - 2) A user may construct a :py:class:`flytekit.common.types.impl.schema.Schema` object (with the correct - column definitions) and pass it as output. - 3) Output can be nulled with a None value. - - From command-line: - Specify a path to the schema object. This path must be accessible from the container when - executing--either by being downloaded from an accessible remote location like s3 or as a local file. - - .. code-block:: python - - @inputs(generic=Types.Schema(), typed=Types.Schema([('a', Types.Integer), ('b', Types.Float)])) - @outputs(b=Types.Schema([('a', Types.Integer), ('b', Types.Float)])) - @python_task - def concat_then_split(wf_params, generic, typed,): - with typed as reader: - # Each chunk is loaded as a pandas.DataFrame object - for df in reader.iter_chunks(): - # Operate on the dataframe - - # Create at a random location specified in flytekit configuration - out = Types.Schema([('a', Types.Integer), ('b', Types.Float)])() - with out as writer: - writer.write( - pandas.DataFrame.from_dict( - { - 'a': [1, 2, 3], - 'b': [5.0, 6.0, 7.0] - } - ) - ) - b.set(out) - """ - - Proto = staticmethod(_proto.create_protobuf) - """ - Proto type wraps a protobuf type to provide interoperability between protobuf and flyte typing system. Using this - type, you can define custom input/output variable types of flyte entities and continue to provide strong typing - syntax. Proto type serializes proto objects as binary (leveraging `flyteidl's Binary literal `_). - Binary serialization of protobufs is the most space-efficient serialization form. Because of the way protobufs are - designed, unmarshalling the serialized proto requires access to the corresponding type. In order to use/visualize - the serialized proto, you will generally need to write custom code in the corresponding component. - - .. note:: - - The protobuf Python library should be installed on the PYTHONPATH to ensure the type engine can access the - appropriate Python code to deserialize the protobuf. - - When used with an SDK-decorated method, expect this behavior from the default type engine: - - As input: - 1) If set, a Python protobuf object of the type specified in the definition. - 2) If not set, a None value. - - As output: - 1) A Python protobuf object matching the type specified by the users. - 2) Set None to null the output. - - From command-line: - A base-64 encoded string of the serialized protobuf. - - .. code-block:: python - - from protos import my_protos_pb2 - - @inputs(a=Types.Proto(my_protos_pb2.Custom)) - @outputs(b=Types.Proto(my_protos_pb2.Custom)) - @python_task - def assert_and_create(wf_params, a, b): - assert a.field1 == 1 - assert a.field2 == 'abc' - b.set( - my_protos_pb2.Custom( - field1=100, - field2='hello' - ) - ) - """ - - GenericProto = staticmethod(_proto.create_generic) - """ - GenericProto type wraps a protobuf type to provide interoperability between protobuf and flyte typing system. Using - this type, you can define custom input/output variable types of flyte entities and continue to provide strong typing - syntax. Proto type serializes proto objects as binary (leveraging `flyteidl's Binary literal `_). - A generic proto is a specialization of the Generic type with added convenience functions to support marshalling/ - unmarshalling of the underlying protobuf object using the protobuf official json marshaller. While GenericProto type - does not produce the most space-efficient representation of protobufs, it's a suitable solution for making protobufs - easily accessible (i.e. humanly readable) in other flyte components (e.g. console, cli... etc.). - - .. note:: - - The protobuf Python library should be installed on the PYTHONPATH to ensure the type engine can access the - appropriate Python code to deserialize the protobuf. - - When used with an SDK-decorated method, expect this behavior from the default type engine: - - As input: - 1) If set, a Python protobuf object of the type specified in the definition. - 2) If not set, a None value. - - As output: - 1) A Python protobuf object matching the type specified by the users. - 2) Set None to null the output. - - From command-line: - A base-64 encoded string of the serialized protobuf. - - .. code-block:: python - - from protos import my_protos_pb2 - - @inputs(a=Types.GenericProto(my_protos_pb2.Custom)) - @outputs(b=Types.GenericProto(my_protos_pb2.Custom)) - @python_task - def assert_and_create(wf_params, a, b): - assert a.field1 == 1 - assert a.field2 == 'abc' - b.set( - my_protos_pb2.Custom( - field1=100, - field2='hello' - ) - ) - """ - - List = staticmethod(_containers.List) - """ - Use this to specify a list of any type--including nested lists. - - When used with an SDK-decorated method, expect this behavior from the default type engine: - - As input: - 1) If set, a Python list populated with values matching the behavior of the list's sub-type. - 2) If not set, a None value. - - As output: - 1) A Python list containing values adhering to the list's sub-type. - 2) Set None to null the output. - - From command-line: - Specify a valid JSON list string. The sub-values will be checked against the sub-type of the list. - - .. note:: - - Shorthand syntax is supported of the form: `[Types.Integer]` in addition to longhand syntax like - `Types.List(Types.Integer)`. Both forms are equivalent. - - .. note:: - - Lists can be arbitrarily deeply nested, however, the typing must be consistent between all sibling values in a - nested list. Syntax for nesting is `[[[Types.Integer]]]` to create a 3D list of integers. - - .. code-block:: python - - @inputs(a=[Types.Integer]) - @outputs(b=[Types.Integer]) - @python_task - def square_each(wf_params, a, b): - b.set( - [x * x for x in a] - ) - """ diff --git a/flytekit/sdk/workflow.py b/flytekit/sdk/workflow.py deleted file mode 100644 index 34561dc592..0000000000 --- a/flytekit/sdk/workflow.py +++ /dev/null @@ -1,132 +0,0 @@ -from typing import Dict - -import six as _six - -import flytekit.common.local_workflow -from flytekit.common import nodes as _nodes -from flytekit.common import promise as _promise -from flytekit.common import workflow as _common_workflow -from flytekit.common.types import helpers as _type_helpers - - -class Input(_promise.Input): - """ - This object should be used to specify inputs. It can be used in conjunction with - :py:meth:`flytekit.common.workflow.workflow` and :py:meth:`flytekit.common.workflow.workflow_class` - """ - - def __init__(self, sdk_type, help=None, **kwargs): - """ - :param flytekit.common.types.base_sdk_types.FlyteSdkType sdk_type: This is the SDK type necessary to create an - input to this workflow. - :param Text help: An optional help string to describe the input to users. - :param bool required: If set, default must be None - :param T default: If this is not a required input, the value will default to this value. Specify as a kwarg. - """ - super(Input, self).__init__("", _type_helpers.python_std_to_sdk_type(sdk_type), help=help, **kwargs) - - -class Output(flytekit.common.local_workflow.Output): - """ - This object should be used to specify outputs. It can be used in conjunction with - :py:meth:`flytekit.common.workflow.workflow` and :py:meth:`flytekit.common.workflow.workflow_class` - """ - - def __init__(self, value, sdk_type=None, help=None): - """ - :param T value: - :param flytekit.common.types.base_sdk_types.FlyteSdkType sdk_type: If specified, the value provided must - match this type exactly. If not provided, the SDK will attempt to infer the type. It is recommended - this value be provided as the SDK might not always be able to infer the correct type. - """ - super(Output, self).__init__( - "", - value, - sdk_type=_type_helpers.python_std_to_sdk_type(sdk_type) if sdk_type else None, - help=help, - ) - - -def workflow_class(_workflow_metaclass=None, on_failure=None, disable_default_launch_plan=False, cls=None): - """ - This is a decorator for wrapping class definitions into workflows. - - .. code-block:: python - - @workflow_class - class MyWorkflow(object): - a = Input(Types.Integer, default=100, help="Tell me something") - b = Input(Types.Float, required=True) - first_task = my_task(a=a) - second_task = my_other_task(b=b, c=first_task.outputs.c) - d = Output(node2.outputs.d) - - - :param T _workflow_metaclass: Do NOT specify this parameter directly. This is the class that is being - wrapped by this decorator. - :param flytekit.models.core.workflow.WorkflowMetadata.OnFailurePolicy on_failure: [Optional] The execution policy - when the workflow detects a failure. - :param bool disable_default_launch_plan: Determines whether to create a default launch plan for the workflow or not. - :param cls: This is the class that will be instantiated from the inputs, outputs, and nodes. This will be used - by users extending the base Flyte programming model. If set, it must be a subclass of - :py:class:`flytekit.common.local_workflow.PythonWorkflow`. - - :rtype: flytekit.common.workflow.SdkWorkflow - """ - - def wrapper(metaclass): - wf = flytekit.common.local_workflow.build_sdk_workflow_from_metaclass( - metaclass, on_failure=on_failure, disable_default_launch_plan=disable_default_launch_plan, cls=cls - ) - return wf - - if _workflow_metaclass is not None: - return wrapper(_workflow_metaclass) - return wrapper - - -def workflow(nodes: Dict[str, _nodes.SdkNode], inputs=None, outputs=None, cls=None, on_failure=None): - """ - This function provides a user-friendly interface for authoring workflows. - - .. code-block:: python - - input_a = Input(Types.Integer, default=100, help="Tell me something") - input_b = Input(Types.Float, required=True) - - node1 = my_task(a=input_a) - node2 = my_other_task(b=input_b, c=node1.outputs.c) - - MyWorkflow = workflow( - workflow_id='my_workflow_id', - inputs={ - 'a': input_a, - 'b': input_b - }, - outputs={ - 'd': Output(node2.outputs.d, sdk_type=Types.Integer, help='This is an integer output') - }, - nodes=[ - node1, - node2 - ] - ) - - :param dict[Text,flytekit.common.nodes.SdkNode] nodes: A list of nodes to put inside the workflow. - :param dict[Text,Input] inputs: [Optional] A dictionary of input descriptors for the workflow. - :param dict[Text,Output] outputs: [Optional] A dictionary of output descriptors for a workflow. - :param T cls: This is the class that will be instantiated from the inputs, outputs, and nodes. This will be used - by users extending the base Flyte programming model. If set, it must be a subclass of - :py:class:`flytekit.common.local_workflow.PythonWorkflow`. - :param flytekit.models.core.workflow.WorkflowMetadata.OnFailurePolicy on_failure: [Optional] The execution policy when the workflow detects a failure. - - :rtype: flytekit.common.local_workflow.SdkRunnableWorkflow - """ - # TODO: Why does Pycharm complain about nodes? - wf = (cls or flytekit.common.local_workflow.SdkRunnableWorkflow).construct_from_class_definition( - inputs=[v.rename_and_return_reference(k) for k, v in sorted(_six.iteritems(inputs or {}))], - outputs=[v.rename_and_return_reference(k) for k, v in sorted(_six.iteritems(outputs or {}))], - nodes=[v.assign_id_and_return(k) for k, v in sorted(_six.iteritems(nodes))], - metadata=_common_workflow._workflow_models.WorkflowMetadata(on_failure=on_failure), - ) - return wf diff --git a/flytekit/testing/__init__.py b/flytekit/testing/__init__.py index 15340a7014..bb75358198 100644 --- a/flytekit/testing/__init__.py +++ b/flytekit/testing/__init__.py @@ -16,5 +16,4 @@ """ -from flytekit.common.tasks.sdk_runnable import SecretsManager from flytekit.core.testing import patch, task_mock diff --git a/flytekit/tools/fast_registration.py b/flytekit/tools/fast_registration.py index f53a9b0b96..f9b9fc6665 100644 --- a/flytekit/tools/fast_registration.py +++ b/flytekit/tools/fast_registration.py @@ -6,11 +6,12 @@ import checksumdir -from flytekit.interfaces.data import data_proxy as _data_proxy -from flytekit.interfaces.data.data_proxy import Data as _Data +from flytekit.core.context_manager import FlyteContextManager _tmp_versions_dir = "tmp/versions" +file_access = FlyteContextManager.current_context().file_access + def compute_digest(source_dir: _os.PathLike) -> str: """ @@ -70,7 +71,7 @@ def upload_package(source_dir: _os.PathLike, identifier: str, remote_location: s print("Local marker for identifier {} already exists, skipping upload".format(identifier)) return full_remote_path - if _Data.data_exists(full_remote_path): + if file_access.exists(full_remote_path): print("Remote file {} already exists, skipping upload".format(full_remote_path)) _write_marker(marker) return full_remote_path @@ -82,7 +83,7 @@ def upload_package(source_dir: _os.PathLike, identifier: str, remote_location: s if dry_run: print("Would upload {} to {}".format(fp.name, full_remote_path)) else: - _Data.put_data(fp.name, full_remote_path) + file_access.put_data(fp.name, full_remote_path) print("Uploaded {} to {}".format(fp.name, full_remote_path)) # Finally, touch the marker file so we have a flag in the future to avoid re-uploading the package dir as an @@ -97,7 +98,7 @@ def download_distribution(additional_distribution: str, destination: str): :param Text additional_distribution: :param _os.PathLike destination: """ - _data_proxy.Data.get_data(additional_distribution, destination) + file_access.get_data(additional_distribution, destination) tarfile_name = _os.path.basename(additional_distribution) file_suffix = _Path(tarfile_name).suffixes if len(file_suffix) != 2 or file_suffix[0] != ".tar" or file_suffix[1] != ".gz": diff --git a/flytekit/tools/lazy_loader.py b/flytekit/tools/lazy_loader.py deleted file mode 100644 index a174769673..0000000000 --- a/flytekit/tools/lazy_loader.py +++ /dev/null @@ -1,118 +0,0 @@ -from __future__ import annotations - -import importlib as _importlib -import sys as _sys -import types as _types -from typing import List - - -class LazyLoadPlugin(object): - LAZY_LOADING_PLUGINS = {} - - def __init__(self, plugin_name, plugin_requirements, related_modules: List[_LazyLoadModule]): - """ - :param Text plugin_name: - :param list[Text] plugin_requirements: - :param list[LazyLoadModule] related_modules: - """ - type(self).LAZY_LOADING_PLUGINS[plugin_name] = plugin_requirements - for m in related_modules: - type(m).tag_with_plugin(plugin_name) - - @classmethod - def get_extras_require(cls): - """ - :rtype: dict[Text,list[Text]] - """ - d = cls.LAZY_LOADING_PLUGINS.copy() - all_plugins_spark2 = [] - all_plugins_spark3 = [] - for k in d: - # Default to Spark 2.4.x in all-spark2 and Spark 3.x in all-spark3. - if k != "spark3": - all_plugins_spark2.extend(d[k]) - if k != "spark": - all_plugins_spark3.extend(d[k]) - - d["all-spark2.4"] = all_plugins_spark2 - d["all-spark3"] = all_plugins_spark3 - # all points to Spark 3.x. - # Spark 2.4 to be fully removed in a future release. - d["all"] = all_plugins_spark3 - return d - - -def lazy_load_module(module: str) -> _types.ModuleType: - """ - :param Text module: - :rtype: _types.ModuleType - """ - - class LazyLoadModule(_LazyLoadModule): - _module = module - _lazy_submodules = dict() - _plugins = [] - - return LazyLoadModule(module) - - -class _LazyLoadModule(_types.ModuleType): - _ERROR_MSG_FMT = ( - "Attempting to use a plugin functionality that requires module " - "`{module}`, but it couldn't be loaded. Please pip install at least one of {plugins} or " - "`flytekit[all]` to get these dependencies.\n" - "\n" - "Original message: {msg}" - ) - - @classmethod - def _load(cls): - module = _sys.modules.get(cls._module) - if not module: - try: - module = _importlib.import_module(cls._module) - except ImportError as e: - raise ImportError(cls._ERROR_MSG_FMT.format(module=cls._module, plugins=cls._plugins, msg=e)) - return module - - def __getattribute__(self, item): - if item in type(self)._lazy_submodules: - return type(self)._lazy_submodules[item] - m = type(self)._load() - return getattr(m, item) - - def __setattr__(self, key, value): - m = type(self)._load() - return setattr(m, key, value) - - @classmethod - def _add_sub_module(cls, submodule): - """ - Add a submodule. - :param Text submodule: This should be a single submodule. Do NOT include periods - :rtype: LazyLoadModule - """ - m = cls._lazy_submodules.get(submodule) - if not m: - m = cls._lazy_submodules[submodule] = lazy_load_module("{}.{}".format(cls._module, submodule)) - return m - - @classmethod - def add_sub_module(cls, submodule): - """ - Add a submodule. - :param Text submodule: If periods are included, it will be added recursively - :rtype: LazyLoadModule - """ - parts = submodule.split(".", 1) - m = cls._add_sub_module(parts[0]) - if len(parts) > 1: - m = type(m).add_sub_module(parts[1]) - return m - - @classmethod - def tag_with_plugin(cls, p: LazyLoadPlugin): - """ - :param LazyLoadPlugin p: - """ - cls._plugins.append(p) diff --git a/flytekit/tools/module_loader.py b/flytekit/tools/module_loader.py index 182b11de72..bc0c46bbbf 100644 --- a/flytekit/tools/module_loader.py +++ b/flytekit/tools/module_loader.py @@ -5,10 +5,6 @@ import sys from typing import Any, Iterator, List, Union -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.local_workflow import SdkRunnableWorkflow as _SdkRunnableWorkflow -from flytekit.common.mixins import registerable as _registerable - def iterate_modules(pkgs): for package_name in pkgs: @@ -56,71 +52,6 @@ def load_workflow_modules(pkgs): pass -def _topo_sort_helper( - obj, - entity_to_module_key, - visited, - recursion_set, - recursion_stack, - include_entities, - ignore_entities, - detect_unreferenced_entities, -): - visited.add(obj) - recursion_stack.append(obj) - if obj in recursion_set: - raise _user_exceptions.FlyteAssertion( - "A cyclical dependency was detected during topological sort of entities. " - "Cycle path was:\n\n\t{}".format("\n\t".join(p for p in recursion_stack[recursion_set[obj] :])) - ) - recursion_set[obj] = len(recursion_stack) - 1 - - if isinstance(obj, _registerable.HasDependencies): - for upstream in obj.upstream_entities: - if upstream.has_registered: - continue - if upstream not in visited: - for m1, k1, o1 in _topo_sort_helper( - upstream, - entity_to_module_key, - visited, - recursion_set, - recursion_stack, - include_entities, - ignore_entities, - detect_unreferenced_entities, - ): - if not o1.has_registered: - yield m1, k1, o1 - - recursion_stack.pop() - del recursion_set[obj] - - if isinstance(obj, include_entities) or not isinstance(obj, ignore_entities): - if obj in entity_to_module_key: - yield entity_to_module_key[obj] + (obj,) - elif detect_unreferenced_entities: - raise _user_exceptions.FlyteAssertion( - f"An entity ({obj.id}) was not found in modules accessible from the workflow packages configuration. Please " - f"ensure that entities in '{obj.instantiated_in}' are moved to a configured packaged, or adjust the configuration." - ) - - -def _get_entity_to_module(pkgs): - entity_to_module_key = {} - for m in iterate_modules(pkgs): - for k in dir(m): - o = m.__dict__[k] - if isinstance(o, _registerable.RegisterableEntity) and not o.has_registered: - if o.instantiated_in == m.__name__: - entity_to_module_key[o] = (m, k) - if isinstance(o, _SdkRunnableWorkflow) and o.should_create_default_launch_plan: - # SDK should create a default launch plan for a workflow. This is a special-case to simplify - # authoring of workflows. - entity_to_module_key[o.create_launch_plan()] = (m, k) - return entity_to_module_key - - def load_module_object_for_type(pkgs, t, additional_path=None): def iterate(): entity_to_module_key = {} @@ -138,66 +69,33 @@ def iterate(): return iterate() -def iterate_registerable_entities_in_order( +def load_object_from_module(object_location: str) -> Any: + """ + # TODO: Handle corner cases, like where the first part is [] maybe + """ + class_obj = object_location.split(".") + class_obj_mod = class_obj[:-1] # e.g. ['flytekit', 'core', 'python_auto_container'] + class_obj_key = class_obj[-1] # e.g. 'default_task_class_obj' + class_obj_mod = importlib.import_module(".".join(class_obj_mod)) + return getattr(class_obj_mod, class_obj_key) + + +def trigger_loading( pkgs, local_source_root=None, - ignore_entities=None, - include_entities=None, - detect_unreferenced_entities=True, ): """ This function will iterate all discovered entities in the given package list. It will then attempt to topologically sort such that any entity with a dependency on another comes later in the list. Note that workflows can reference other workflows and launch plans. + :param list[Text] pkgs: :param Text local_source_root: - :param set[type] ignore_entities: If specified, ignore these entities while doing a topological sort. All other - entities will be taken. Only one of ignore_entities or include_entities can be set. - :param set[type] include_entities: If specified, include these entities while doing a topological sort. All - other entities will be ignored. Only one of ignore_entities or include_entities can be set. - :param bool detect_unreferenced_entities: If true, we will raise exceptions on entities not included in the package - configuration. - :rtype: module, Text, flytekit.common.mixins.registerable.RegisterableEntity """ - if ignore_entities and include_entities: - raise _user_exceptions.FlyteAssertion("ignore_entities and include_entities cannot both be set") - elif not ignore_entities and not include_entities: - include_entities = (object,) - ignore_entities = tuple() - else: - ignore_entities = tuple(list(ignore_entities or set([object]))) - include_entities = tuple(list(include_entities or set())) - if local_source_root is not None: with add_sys_path(local_source_root): - entity_to_module_key = _get_entity_to_module(pkgs) + for _ in iterate_modules(pkgs): + ... else: - entity_to_module_key = _get_entity_to_module(pkgs) - - visited = set() - for o in entity_to_module_key.keys(): - if o not in visited: - recursion_set = dict() - recursion_stack = [] - for m, k, o2 in _topo_sort_helper( - o, - entity_to_module_key, - visited, - recursion_set, - recursion_stack, - include_entities, - ignore_entities, - detect_unreferenced_entities=detect_unreferenced_entities, - ): - yield m, k, o2 - - -def load_object_from_module(object_location: str) -> Any: - """ - # TODO: Handle corner cases, like where the first part is [] maybe - """ - class_obj = object_location.split(".") - class_obj_mod = class_obj[:-1] # e.g. ['flytekit', 'core', 'python_auto_container'] - class_obj_key = class_obj[-1] # e.g. 'default_task_class_obj' - class_obj_mod = importlib.import_module(".".join(class_obj_mod)) - return getattr(class_obj_mod, class_obj_key) + for _ in iterate_modules(pkgs): + ... diff --git a/flytekit/common/translator.py b/flytekit/tools/translator.py similarity index 99% rename from flytekit/common/translator.py rename to flytekit/tools/translator.py index deca82f7fd..83d724e4ce 100644 --- a/flytekit/common/translator.py +++ b/flytekit/tools/translator.py @@ -1,8 +1,7 @@ from collections import OrderedDict from typing import Callable, Dict, List, Optional, Tuple, Union -from flytekit.common import constants as _common_constants -from flytekit.common.utils import _dnsify +from flytekit.core import constants as _common_constants from flytekit.core.base_task import PythonTask from flytekit.core.condition import BranchNode from flytekit.core.context_manager import SerializationSettings @@ -11,6 +10,7 @@ from flytekit.core.python_auto_container import PythonAutoContainerTask from flytekit.core.reference_entity import ReferenceEntity, ReferenceSpec, ReferenceTemplate from flytekit.core.task import ReferenceTask +from flytekit.core.utils import _dnsify from flytekit.core.workflow import ReferenceWorkflow, WorkflowBase from flytekit.models import common as _common_models from flytekit.models import interface as interface_models diff --git a/flytekit/type_engines/__init__.py b/flytekit/type_engines/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/flytekit/type_engines/common.py b/flytekit/type_engines/common.py deleted file mode 100644 index 33927325ce..0000000000 --- a/flytekit/type_engines/common.py +++ /dev/null @@ -1,31 +0,0 @@ -import abc as _abc - - -class TypeEngine(object, metaclass=_abc.ABCMeta): - @_abc.abstractmethod - def python_std_to_sdk_type(self, t): - """ - Converts a standard format for specifying types in Python to the Flyte typing structure. - :param T t: User input. Usually of the form: Types.Integer, [Types.Integer], {Types.String: - Types.Integer}, etc. - :rtype: flytekit.common.types.base_sdk_types.FlyteSdkType - """ - pass - - @_abc.abstractmethod - def get_sdk_type_from_literal_type(self, literal_type): - """ - Takes the Flyte spec language and converts to an SDK object. - :param flytekit.models.types.LiteralType literal_type: - :rtype: flytekit.common.types.base_sdk_types.FlyteSdkType - """ - pass - - @_abc.abstractmethod - def infer_sdk_type_from_literal(self, literal): - """ - From a literal value, we infer the correct SDK type. - :param flytekit.models.literals.Literal literal: - :rtype: flytekit.common.types.base_sdk_types.FlyteSdkType - """ - pass diff --git a/flytekit/type_engines/default/__init__.py b/flytekit/type_engines/default/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/flytekit/type_engines/default/flyte.py b/flytekit/type_engines/default/flyte.py deleted file mode 100644 index 0ec3a4c982..0000000000 --- a/flytekit/type_engines/default/flyte.py +++ /dev/null @@ -1,193 +0,0 @@ -import importlib as _importer -from typing import Type - -from flytekit.common.exceptions import system as _system_exceptions -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.types import base_sdk_types as _base_sdk_types -from flytekit.common.types import blobs as _blobs -from flytekit.common.types import containers as _container_types -from flytekit.common.types import helpers as _helpers -from flytekit.common.types import primitives as _primitive_types -from flytekit.common.types import proto as _proto -from flytekit.common.types import schema as _schema -from flytekit.models import types as _literal_type_models -from flytekit.models.core import types as _core_types - - -def _load_type_from_tag(tag: str) -> Type: - """ - Loads python type from tag - """ - - if "." not in tag: - raise _user_exceptions.FlyteValueException( - tag, - "Protobuf tag must include at least one '.' to delineate package and object name.", - ) - - module, name = tag.rsplit(".", 1) - try: - pb_module = _importer.import_module(module) - except ImportError: - raise _user_exceptions.FlyteAssertion( - "Could not resolve the protobuf definition @ {}. Is the protobuf library installed?".format(module) - ) - - if not hasattr(pb_module, name): - raise _user_exceptions.FlyteAssertion("Could not find the protobuf named: {} @ {}.".format(name, module)) - - return getattr(pb_module, name) - - -def _proto_sdk_type_from_tag(tag): - """ - :param Text tag: - :rtype: _proto.Protobuf - """ - return _proto.create_protobuf(_load_type_from_tag(tag)) - - -def _generic_proto_sdk_type_from_tag(tag: str) -> Type[_proto.GenericProtobuf]: - """ - :param Text tag: - :rtype: _proto.GenericProtobuf - """ - - return _proto.create_generic(_load_type_from_tag(tag)) - - -class FlyteDefaultTypeEngine(object): - _SIMPLE_TYPE_LOOKUP_TABLE = { - _literal_type_models.SimpleType.INTEGER: _primitive_types.Integer, - _literal_type_models.SimpleType.FLOAT: _primitive_types.Float, - _literal_type_models.SimpleType.BOOLEAN: _primitive_types.Boolean, - _literal_type_models.SimpleType.DATETIME: _primitive_types.Datetime, - _literal_type_models.SimpleType.DURATION: _primitive_types.Timedelta, - _literal_type_models.SimpleType.NONE: _base_sdk_types.Void, - _literal_type_models.SimpleType.STRING: _primitive_types.String, - _literal_type_models.SimpleType.STRUCT: _primitive_types.Generic, - } - - def python_std_to_sdk_type(self, t): - """ - :param T t: User input. Should be of the form: Types.Integer, [Types.Integer], {Types.String: - Types.Integer}, etc. - :rtype: flytekit.common.types.base_sdk_types.FlyteSdkType - """ - if isinstance(t, list): - if len(t) != 1: - raise _user_exceptions.FlyteAssertion( - "When specifying a list type, there must be exactly one element in " - "the list describing the contained type." - ) - return _container_types.List(_helpers.python_std_to_sdk_type(t[0])) - elif isinstance(t, dict): - raise _user_exceptions.FlyteAssertion("Map types are not yet implemented.") - elif isinstance(t, _base_sdk_types.FlyteSdkType): - return t - else: - raise _user_exceptions.FlyteTypeException( - type(t), - _base_sdk_types.FlyteSdkType, - additional_msg="Should be of form similar to: Types.Integer, [Types.Integer], {Types.String: " - "Types.Integer}", - received_value=t, - ) - - def get_sdk_type_from_literal_type(self, literal_type): - """ - :param flytekit.models.types.LiteralType literal_type: - :rtype: flytekit.common.types.base_sdk_types.FlyteSdkType - """ - if literal_type.collection_type is not None: - return _container_types.List(_helpers.get_sdk_type_from_literal_type(literal_type.collection_type)) - elif literal_type.map_value_type is not None: - raise NotImplementedError("TODO: Implement map") - elif literal_type.schema is not None: - return _schema.schema_instantiator_from_proto(literal_type.schema) - elif literal_type.blob is not None: - return self._get_blob_impl_from_type(literal_type.blob) - elif literal_type.simple is not None: - if ( - literal_type.simple == _literal_type_models.SimpleType.BINARY - and _proto.Protobuf.PB_FIELD_KEY in literal_type.metadata - ): - return _proto_sdk_type_from_tag(literal_type.metadata[_proto.Protobuf.PB_FIELD_KEY]) - if ( - literal_type.simple == _literal_type_models.SimpleType.STRUCT - and literal_type.metadata - and _proto.Protobuf.PB_FIELD_KEY in literal_type.metadata - ): - return _generic_proto_sdk_type_from_tag(literal_type.metadata[_proto.Protobuf.PB_FIELD_KEY]) - sdk_type = self._SIMPLE_TYPE_LOOKUP_TABLE.get(literal_type.simple) - if sdk_type is None: - raise NotImplementedError( - "We haven't implemented this type yet: Simple type={}".format(literal_type.simple) - ) - return sdk_type - else: - raise _system_exceptions.FlyteSystemAssertion( - "An unrecognized literal type was received: {}".format(literal_type) - ) - - def infer_sdk_type_from_literal(self, literal): # noqa - """ - :param flytekit.models.literals.Literal literal: - :rtype: flytekit.common.types.base_sdk_types.FlyteSdkType - """ - if literal.collection is not None: - if len(literal.collection.literals) > 0: - sdk_type = _container_types.List(_helpers.infer_sdk_type_from_literal(literal.collection.literals[0])) - else: - sdk_type = _container_types.List(_base_sdk_types.Void) - elif literal.map is not None: - raise NotImplementedError("TODO: Implement map") - elif literal.scalar.blob is not None: - sdk_type = self._get_blob_impl_from_type(literal.scalar.blob.metadata.type) - elif literal.scalar.none_type is not None: - sdk_type = _base_sdk_types.Void - elif literal.scalar.schema is not None: - sdk_type = _schema.schema_instantiator_from_proto(literal.scalar.schema.type) - elif literal.scalar.error is not None: - raise NotImplementedError("TODO: Implement error from literal map") - elif literal.scalar.generic is not None: - sdk_type = _primitive_types.Generic - elif literal.scalar.binary is not None: - if literal.scalar.binary.tag.startswith(_proto.Protobuf.TAG_PREFIX): - sdk_type = _proto_sdk_type_from_tag(literal.scalar.binary.tag[len(_proto.Protobuf.TAG_PREFIX) :]) - else: - raise NotImplementedError("TODO: Binary is only supported for protobuf types currently") - elif literal.scalar.primitive.boolean is not None: - sdk_type = _primitive_types.Boolean - elif literal.scalar.primitive.datetime is not None: - sdk_type = _primitive_types.Datetime - elif literal.scalar.primitive.duration is not None: - sdk_type = _primitive_types.Timedelta - elif literal.scalar.primitive.float_value is not None: - sdk_type = _primitive_types.Float - elif literal.scalar.primitive.integer is not None: - sdk_type = _primitive_types.Integer - elif literal.scalar.primitive.string_value is not None: - sdk_type = _primitive_types.String - else: - raise _system_exceptions.FlyteSystemAssertion("Received unknown literal: {}".format(literal)) - return sdk_type - - def _get_blob_impl_from_type(self, blob_type): - """ - :param flytekit.models.core.types.BlobType blob_type: - :rtype: flytekit.common.types.base_sdk_types.FlyteSdkType - """ - if blob_type.dimensionality == _core_types.BlobType.BlobDimensionality.SINGLE: - if blob_type.format == "csv": - return _blobs.CSV - else: - return _blobs.Blob - elif blob_type.dimensionality == _core_types.BlobType.BlobDimensionality.MULTIPART: - if blob_type.format == "csv": - return _blobs.MultiPartCSV - else: - return _blobs.MultiPartBlob - raise _system_exceptions.FlyteSystemAssertion( - "Flyte's base type engine does not support this type of blob. Value: {}".format(blob_type) - ) diff --git a/flytekit/types/schema/types.py b/flytekit/types/schema/types.py index b9d2a51372..0afed18158 100644 --- a/flytekit/types/schema/types.py +++ b/flytekit/types/schema/types.py @@ -10,6 +10,7 @@ from typing import Type import numpy as _np +import pandas from dataclasses_json import config, dataclass_json from marshmallow import fields @@ -17,7 +18,6 @@ from flytekit.core.type_engine import T, TypeEngine, TypeTransformer, TypeTransformerFailedError from flytekit.models.literals import Literal, Scalar, Schema from flytekit.models.types import LiteralType, SchemaType -from flytekit.plugins import pandas T = typing.TypeVar("T") diff --git a/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker/hpo.py b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker/hpo.py index 629c8c588d..ae805d0e81 100644 --- a/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker/hpo.py +++ b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker/hpo.py @@ -9,10 +9,9 @@ from google.protobuf.json_format import MessageToDict from flytekit import FlyteContext -from flytekit.common.types import primitives from flytekit.extend import DictTransformer, PythonTask, SerializationSettings, TypeEngine, TypeTransformer from flytekit.models.literals import Literal -from flytekit.models.types import LiteralType +from flytekit.models.types import LiteralType, SimpleType from .models import hpo_job as _hpo_job_model from .models import parameter_ranges as _params @@ -107,7 +106,7 @@ def __init__(self): super().__init__("sagemaker-hpojobconfig-transformer", _hpo_job_model.HyperparameterTuningJobConfig) def get_literal_type(self, t: Type[_hpo_job_model.HyperparameterTuningJobConfig]) -> LiteralType: - return primitives.Generic.to_flyte_literal_type() + return LiteralType(simple=SimpleType.STRUCT, metadata=None) def to_literal( self, @@ -139,7 +138,7 @@ def __init__(self): super().__init__("sagemaker-paramrange-transformer", _params.ParameterRangeOneOf) def get_literal_type(self, t: Type[_params.ParameterRangeOneOf]) -> LiteralType: - return primitives.Generic.to_flyte_literal_type() + return LiteralType(simple=SimpleType.STRUCT, metadata=None) def to_literal( self, diff --git a/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker/models/parameter_ranges.py b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker/models/parameter_ranges.py index 4328749d72..0df8f42dba 100644 --- a/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker/models/parameter_ranges.py +++ b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker/models/parameter_ranges.py @@ -2,7 +2,7 @@ from flyteidl.plugins.sagemaker import parameter_ranges_pb2 as _idl_parameter_ranges -from flytekit.common.exceptions import user +from flytekit.exceptions import user from flytekit.models import common as _common diff --git a/plugins/flytekit-aws-sagemaker/tests/test_hpo.py b/plugins/flytekit-aws-sagemaker/tests/test_hpo.py index 28a226e696..e52994c664 100644 --- a/plugins/flytekit-aws-sagemaker/tests/test_hpo.py +++ b/plugins/flytekit-aws-sagemaker/tests/test_hpo.py @@ -24,7 +24,7 @@ from flytekitplugins.awssagemaker.training import SagemakerBuiltinAlgorithmsTask, SagemakerTrainingJobConfig from flytekit import FlyteContext -from flytekit.common.types.primitives import Generic +from flytekit.models.types import LiteralType, SimpleType from .test_training import _get_reg_settings @@ -93,7 +93,7 @@ def test_hpo_for_builtin(): def test_hpoconfig_transformer(): t = HPOTuningJobConfigTransformer() - assert t.get_literal_type(HyperparameterTuningJobConfig) == Generic.to_flyte_literal_type() + assert t.get_literal_type(HyperparameterTuningJobConfig) == LiteralType(simple=SimpleType.STRUCT) o = HyperparameterTuningJobConfig( tuning_strategy=1, tuning_objective=HyperparameterTuningObjective( @@ -113,7 +113,7 @@ def test_hpoconfig_transformer(): def test_parameter_ranges_transformer(): t = ParameterRangesTransformer() - assert t.get_literal_type(ParameterRangeOneOf) == Generic.to_flyte_literal_type() + assert t.get_literal_type(ParameterRangeOneOf) == LiteralType(simple=SimpleType.STRUCT) o = ParameterRangeOneOf(param=IntegerParameterRange(10, 0, 1)) ctx = FlyteContext.current_context() lit = t.to_literal(ctx, python_val=o, python_type=ParameterRangeOneOf, expected=None) diff --git a/plugins/flytekit-aws-sagemaker/tests/test_training.py b/plugins/flytekit-aws-sagemaker/tests/test_training.py index 3e6ada1ff5..a48d3c9f39 100644 --- a/plugins/flytekit-aws-sagemaker/tests/test_training.py +++ b/plugins/flytekit-aws-sagemaker/tests/test_training.py @@ -14,7 +14,7 @@ import flytekit from flytekit import task -from flytekit.common.tasks.sdk_runnable import ExecutionParameters +from flytekit.core.context_manager import ExecutionParameters from flytekit.extend import Image, ImageConfig, SerializationSettings diff --git a/plugins/flytekit-k8s-pod/flytekitplugins/pod/task.py b/plugins/flytekit-k8s-pod/flytekitplugins/pod/task.py index 05b7cdd5e1..6cbf7d57fa 100644 --- a/plugins/flytekit-k8s-pod/flytekitplugins/pod/task.py +++ b/plugins/flytekit-k8s-pod/flytekitplugins/pod/task.py @@ -5,7 +5,7 @@ from kubernetes.client.models import V1Container, V1EnvVar, V1PodSpec, V1ResourceRequirements from flytekit import FlyteContext, PythonFunctionTask -from flytekit.common.exceptions import user as _user_exceptions +from flytekit.exceptions import user as _user_exceptions from flytekit.extend import Promise, SerializationSettings, TaskPlugins from flytekit.models import task as _task_models diff --git a/plugins/flytekit-k8s-pod/tests/test_pod.py b/plugins/flytekit-k8s-pod/tests/test_pod.py index ed0229f02b..82a0bcf6f8 100644 --- a/plugins/flytekit-k8s-pod/tests/test_pod.py +++ b/plugins/flytekit-k8s-pod/tests/test_pod.py @@ -8,10 +8,10 @@ from kubernetes.client.models import V1Container, V1EnvVar, V1PodSpec, V1ResourceRequirements, V1VolumeMount from flytekit import Resources, TaskMetadata, dynamic, map_task, task -from flytekit.common.translator import get_serializable from flytekit.core import context_manager from flytekit.core.context_manager import FastSerializationSettings from flytekit.extend import ExecutionState, Image, ImageConfig, SerializationSettings +from flytekit.tools.translator import get_serializable def get_pod_spec(): diff --git a/plugins/flytekit-papermill/flytekitplugins/papermill/task.py b/plugins/flytekit-papermill/flytekitplugins/papermill/task.py index 1a7845a8fe..4b1e183952 100644 --- a/plugins/flytekit-papermill/flytekitplugins/papermill/task.py +++ b/plugins/flytekit-papermill/flytekitplugins/papermill/task.py @@ -11,7 +11,7 @@ from nbconvert import HTMLExporter from flytekit import FlyteContext, PythonInstanceTask -from flytekit.common.tasks.sdk_runnable import ExecutionParameters +from flytekit.core.context_manager import ExecutionParameters from flytekit.extend import Interface, TaskPlugins, TypeEngine from flytekit.models.literals import LiteralMap from flytekit.types.file import HTMLPage, PythonNotebook diff --git a/plugins/flytekit-spark/flytekitplugins/spark/models.py b/plugins/flytekit-spark/flytekitplugins/spark/models.py index d03949f1ac..53b1620331 100644 --- a/plugins/flytekit-spark/flytekitplugins/spark/models.py +++ b/plugins/flytekit-spark/flytekitplugins/spark/models.py @@ -1,10 +1,17 @@ +import enum import typing from flyteidl.plugins import spark_pb2 as _spark_task -from flytekit.common.exceptions import user as _user_exceptions +from flytekit.exceptions import user as _user_exceptions from flytekit.models import common as _common -from flytekit.sdk.spark_types import SparkType as _spark_type + + +class SparkType(enum.Enum): + PYTHON = 1 + SCALA = 2 + JAVA = 3 + R = 4 class SparkJob(_common.FlyteIdlEntity): @@ -102,13 +109,13 @@ def to_flyte_idl(self): :rtype: flyteidl.plugins.spark_pb2.SparkJob """ - if self.spark_type == _spark_type.PYTHON: + if self.spark_type == SparkType.PYTHON: application_type = _spark_task.SparkApplication.PYTHON - elif self.spark_type == _spark_type.JAVA: + elif self.spark_type == SparkType.JAVA: application_type = _spark_task.SparkApplication.JAVA - elif self.spark_type == _spark_type.SCALA: + elif self.spark_type == SparkType.SCALA: application_type = _spark_task.SparkApplication.SCALA - elif self.spark_type == _spark_type.R: + elif self.spark_type == SparkType.R: application_type = _spark_task.SparkApplication.R else: raise _user_exceptions.FlyteValidationException("Invalid Spark Application Type Specified") @@ -129,13 +136,13 @@ def from_flyte_idl(cls, pb2_object): :rtype: SparkJob """ - application_type = _spark_type.PYTHON + application_type = SparkType.PYTHON if pb2_object.type == _spark_task.SparkApplication.JAVA: - application_type = _spark_type.JAVA + application_type = SparkType.JAVA elif pb2_object.type == _spark_task.SparkApplication.SCALA: - application_type = _spark_type.SCALA + application_type = SparkType.SCALA elif pb2_object.type == _spark_task.SparkApplication.R: - application_type = _spark_type.R + application_type = SparkType.R return cls( type=application_type, diff --git a/plugins/flytekit-spark/flytekitplugins/spark/task.py b/plugins/flytekit-spark/flytekitplugins/spark/task.py index bb20da8b48..37ae03e913 100644 --- a/plugins/flytekit-spark/flytekitplugins/spark/task.py +++ b/plugins/flytekit-spark/flytekitplugins/spark/task.py @@ -7,11 +7,10 @@ from pyspark.sql import SparkSession from flytekit import FlyteContextManager, PythonFunctionTask -from flytekit.common.tasks.sdk_runnable import ExecutionParameters +from flytekit.core.context_manager import ExecutionParameters from flytekit.extend import ExecutionState, SerializationSettings, TaskPlugins -from flytekit.sdk.spark_types import SparkType -from .models import SparkJob +from .models import SparkJob, SparkType @dataclass diff --git a/plugins/flytekit-spark/tests/test_spark_task.py b/plugins/flytekit-spark/tests/test_spark_task.py index 4fc80d541f..dfda716a11 100644 --- a/plugins/flytekit-spark/tests/test_spark_task.py +++ b/plugins/flytekit-spark/tests/test_spark_task.py @@ -3,7 +3,7 @@ import flytekit from flytekit import task -from flytekit.common.tasks.sdk_runnable import ExecutionParameters +from flytekit.core.context_manager import ExecutionParameters from flytekit.extend import Image, ImageConfig, SerializationSettings diff --git a/plugins/flytekit-sqlalchemy/tests/test_sql_tracker.py b/plugins/flytekit-sqlalchemy/tests/test_sql_tracker.py index 404ef2fb23..93104eabdd 100644 --- a/plugins/flytekit-sqlalchemy/tests/test_sql_tracker.py +++ b/plugins/flytekit-sqlalchemy/tests/test_sql_tracker.py @@ -1,8 +1,8 @@ from collections import OrderedDict -from flytekit.common.translator import get_serializable from flytekit.core import context_manager from flytekit.core.context_manager import Image, ImageConfig +from flytekit.tools.translator import get_serializable from .test_task import tk as not_tk diff --git a/plugins/flytekit-sqlalchemy/tests/test_task.py b/plugins/flytekit-sqlalchemy/tests/test_task.py index 2c853f7811..167c8e796d 100644 --- a/plugins/flytekit-sqlalchemy/tests/test_task.py +++ b/plugins/flytekit-sqlalchemy/tests/test_task.py @@ -11,8 +11,8 @@ from flytekitplugins.sqlalchemy.task import SQLAlchemyTaskExecutor from flytekit import kwtypes, task, workflow +from flytekit.core.context_manager import SecretsManager from flytekit.models.security import Secret -from flytekit.testing import SecretsManager from flytekit.types.schema import FlyteSchema tk = SQLAlchemyTask( diff --git a/setup.py b/setup.py index bb5f9d1f1c..ae028038bb 100644 --- a/setup.py +++ b/setup.py @@ -2,45 +2,16 @@ from setuptools import find_packages, setup # noqa -# from flytekit.tools.lazy_loader import LazyLoadPlugin # noqa -# extras_require = LazyLoadPlugin.get_extras_require() - MIN_PYTHON_VERSION = (3, 7) CURRENT_PYTHON = sys.version_info[:2] -if CURRENT_PYTHON == (3, 6): - print( - f"Flytekit native typed API is supported for python versions {MIN_PYTHON_VERSION}+, Python 3.6 is supported" - f" only for legacy Flytekit API. This will be deprecated when Python 3.6 reaches end of life (Dec 23rd, 2021)," - f" we recommend migrating to the new API" - ) -elif CURRENT_PYTHON < MIN_PYTHON_VERSION: +if CURRENT_PYTHON < MIN_PYTHON_VERSION: print( f"Flytekit API is only supported for Python version is {MIN_PYTHON_VERSION}+. Detected you are on" f" version {CURRENT_PYTHON}, installation will not proceed!" ) sys.exit(-1) -spark = ["pyspark>=2.4.0,<3.0.0"] -spark3 = ["pyspark>=3.0.0"] -sidecar = ["k8s-proto>=0.0.3,<1.0.0"] -schema = ["numpy>=1.14.0,<2.0.0", "pandas>=0.22.0,<2.0.0", "pyarrow>=4.0.0"] -hive_sensor = ["hmsclient>=0.0.1,<1.0.0"] -notebook = ["papermill>=1.2.0", "nbconvert>=6.0.7", "ipykernel>=5.0.0,<6.0.0"] -sagemaker = ["sagemaker-training>=3.6.2,<4.0.0"] - -all_but_spark = sidecar + schema + hive_sensor + notebook + sagemaker - -extras_require = { - "spark": spark, - "spark3": spark3, - "sidecar": sidecar, - "schema": schema, - "hive_sensor": hive_sensor, - "notebook": notebook, - "sagemaker": sagemaker, - "all-spark2.4": spark + all_but_spark, - "all": spark3 + all_but_spark, -} +extras_require = {} __version__ = "0.0.0+develop" @@ -71,7 +42,7 @@ "click>=6.6,<8.0", "croniter>=0.3.20,<4.0.0", "deprecated>=1.0,<2.0", - "python-dateutil<=2.8.1,>=2.1", + "python-dateutil>=2.1", "grpcio>=1.3.0,<2.0", "protobuf>=3.6.1,<4", "python-json-logger>=2.0.0", @@ -80,7 +51,6 @@ "keyring>=18.0.1", "requests>=2.18.4,<3.0.0", "responses>=0.10.7", - "six>=1.9.0,<2.0.0", "sortedcontainers>=1.5.9<3.0.0", "statsd>=3.0.0,<4.0.0", "urllib3>=1.22,<2.0.0", @@ -104,7 +74,7 @@ "flytekit/bin/entrypoint.py", ], license="apache2", - python_requires=">=3.6", + python_requires=">=3.7", classifiers=[ "Intended Audience :: Science/Research", "Intended Audience :: Developers", diff --git a/tests/flytekit/common/parameterizers.py b/tests/flytekit/common/parameterizers.py index 33bd7712ee..2d5504b32c 100644 --- a/tests/flytekit/common/parameterizers.py +++ b/tests/flytekit/common/parameterizers.py @@ -1,10 +1,6 @@ from datetime import timedelta from itertools import product -from six.moves import range - -from flytekit.common.types.impl import blobs as _blob_impl -from flytekit.common.types.impl import schema as _schema_impl from flytekit.models import interface, literals, security, task, types from flytekit.models.core import identifier from flytekit.models.core import types as _core_types @@ -253,39 +249,23 @@ ( literals.Scalar( union=literals.Union( - value=literals.Literal( - scalar=literals.Scalar( - primitive=literals.Primitive( - integer=10 - ) - ) - ), + value=literals.Literal(scalar=literals.Scalar(primitive=literals.Primitive(integer=10))), stored_type=types.LiteralType( - simple=types.SimpleType.INTEGER, - structure=types.TypeStructure(tag="int") - ) + simple=types.SimpleType.INTEGER, structure=types.TypeStructure(tag="int") + ), ) ), - 10 + 10, ), ( literals.Scalar( union=literals.Union( - value=literals.Literal( - scalar=literals.Scalar( - primitive=literals.Primitive( - string_value="test" - ) - ) - ), - stored_type=types.LiteralType( - simple=types.SimpleType.STRING, - structure=types.TypeStructure(tag="str") - ) + value=literals.Literal(scalar=literals.Scalar(primitive=literals.Primitive(string_value="test"))), + stored_type=types.LiteralType(simple=types.SimpleType.STRING, structure=types.TypeStructure(tag="str")), ) ), - "test" - ) + "test", + ), ] LIST_OF_SCALAR_LITERALS_AND_PYTHON_VALUE = [ diff --git a/tests/flytekit/common/task_definitions.py b/tests/flytekit/common/task_definitions.py deleted file mode 100644 index 56534fd7ea..0000000000 --- a/tests/flytekit/common/task_definitions.py +++ /dev/null @@ -1,9 +0,0 @@ -from flytekit.sdk.tasks import inputs, outputs, python_task -from flytekit.sdk.types import Types - - -@inputs(a=Types.Integer) -@outputs(b=Types.Integer) -@python_task -def add_one(wf_params, a, b): - b.set(a + 1) diff --git a/tests/flytekit/common/workflows/__init__.py b/tests/flytekit/common/workflows/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/tests/flytekit/common/workflows/batch.py b/tests/flytekit/common/workflows/batch.py deleted file mode 100644 index d3e095dd82..0000000000 --- a/tests/flytekit/common/workflows/batch.py +++ /dev/null @@ -1,128 +0,0 @@ -from six import moves as _six_moves - -from flytekit.sdk.tasks import dynamic_task, inputs, outputs, python_task -from flytekit.sdk.types import Types -from flytekit.sdk.workflow import Input, Output, workflow_class - - -@outputs(out_ints=[Types.Integer]) -@dynamic_task -def sample_batch_task_sq(wf_params, out_ints): - wf_params.stats.incr("task_run") - res2 = [] - for i in _six_moves.range(0, 3): - task = sq_sub_task(in1=i) - yield task - res2.append(task.outputs.out1) - out_ints.set(res2) - - -@outputs(out_str=[Types.String], out_ints=[[Types.Integer]]) -@dynamic_task -def no_inputs_sample_batch_task(wf_params, out_str, out_ints): - wf_params.stats.incr("task_run") - res = ["I'm the first result"] - for i in _six_moves.range(0, 3): - task = sub_task(in1=i) - yield task - res.append(task.outputs.out1) - res.append("I'm after each sub-task result") - res.append("I'm the last result") - - res2 = [] - for i in _six_moves.range(0, 3): - task = int_sub_task(in1=i) - yield task - res2.append(task.outputs.out1) - - # Nested batch tasks - task = sample_batch_task_sq() - yield task - res2.append(task.outputs.out_ints) - - task = sample_batch_task_sq() - yield task - res2.append(task.outputs.out_ints) - - out_str.set(res) - out_ints.set(res2) - - -@inputs(in1=Types.Integer) -@outputs(out_str=[Types.String]) -@dynamic_task(cache=True, cache_version="1") -def sample_batch_task_beatles_cached(wf_params, in1, out_str): - wf_params.stats.incr("task_run") - res2 = [] - for i in _six_moves.range(0, in1): - task = sample_beatles_lyrics_cached(in1=i) - yield task - res2.append(task.outputs.out1) - out_str.set(res2) - - -@inputs(in1=Types.Integer) -@outputs(out1=Types.String) -@python_task(cache=True, cache_version="1") -def sample_beatles_lyrics_cached(wf_params, in1, out1): - wf_params.stats.incr("task_run") - lyrics = ["Ob-La-Di, Ob-La-Da", "When I'm 64", "Yesterday"] - out1.set(lyrics[in1 % 3]) - - -@inputs(in1=Types.Integer) -@outputs(out1=Types.String) -@python_task -def sub_task(wf_params, in1, out1): - wf_params.stats.incr("task_run") - out1.set("hello {}".format(in1)) - - -@inputs(in1=Types.Integer) -@outputs(out1=[Types.Integer]) -@python_task -def int_sub_task(wf_params, in1, out1): - wf_params.stats.incr("task_run") - out1.set([in1, in1 * 2, in1 * 3]) - - -@inputs(in1=Types.Integer) -@outputs(out1=Types.Integer) -@python_task -def sq_sub_task(wf_params, in1, out1): - wf_params.stats.incr("task_run") - out1.set(in1 * in1) - - -@inputs(ints_to_print=[[Types.Integer]], strings_to_print=[Types.String]) -@python_task(cache_version="1") -def print_every_time(wf_params, ints_to_print, strings_to_print): - wf_params.stats.incr("task_run") - print("Expected Int values: {}".format([[0, 0, 0], [1, 2, 3], [2, 4, 6], [0, 1, 4], [0, 1, 4]])) - print("Actual Int values: {}".format(ints_to_print)) - - print( - "Expected String values: {}".format( - [ - u"I'm the first result", - u"hello 0", - u"I'm after each sub-task result", - u"hello 1", - u"I'm after each sub-task result", - u"hello 2", - u"I'm after each sub-task result", - u"I'm the last result", - ] - ) - ) - print("Actual String values: {}".format(strings_to_print)) - - -@workflow_class -class BatchTasksWorkflow(object): - num_subtasks = Input(Types.Integer, default=3) - task1 = no_inputs_sample_batch_task() - task2 = sample_batch_task_beatles_cached(in1=num_subtasks) - t = print_every_time(ints_to_print=task1.outputs.out_ints, strings_to_print=task1.outputs.out_str) - ints_out = Output(task1.outputs.out_ints, sdk_type=[[Types.Integer]]) - str_out = Output(task2.outputs.out_str, sdk_type=[Types.String]) diff --git a/tests/flytekit/common/workflows/dynamic_workflows.py b/tests/flytekit/common/workflows/dynamic_workflows.py deleted file mode 100644 index fabab6367b..0000000000 --- a/tests/flytekit/common/workflows/dynamic_workflows.py +++ /dev/null @@ -1,39 +0,0 @@ -from flytekit.sdk import tasks as _tasks -from flytekit.sdk import workflow as _workflow -from flytekit.sdk.types import Types as _Types -from flytekit.sdk.workflow import Input, Output, workflow_class - - -@_tasks.inputs(num=_Types.Integer) -@_tasks.outputs(out=_Types.Integer) -@_tasks.python_task -def inner_task(wf_params, num, out): - wf_params.logging.info("Running inner task... setting output to input") - out.set(num) - - -@_workflow.workflow_class() -class IdentityWorkflow(object): - a = _workflow.Input(_Types.Integer, default=5, help="Input for inner workflow") - odd_nums_task = inner_task(num=a) - task_output = _workflow.Output(odd_nums_task.outputs.out, sdk_type=_Types.Integer) - - -id_lp = IdentityWorkflow.create_launch_plan() - - -@_tasks.inputs(num=_Types.Integer) -@_tasks.outputs(out=_Types.Integer) -@_tasks.dynamic_task -def lp_yield_task(wf_params, num, out): - wf_params.logging.info("Running inner task... yielding a launchplan") - identity_lp_execution = id_lp(a=num) - yield identity_lp_execution - out.set(identity_lp_execution.outputs.task_output) - - -@workflow_class -class DynamicLaunchPlanCaller(object): - outer_a = Input(_Types.Integer, default=5, help="Input for inner workflow") - lp_task = lp_yield_task(num=outer_a) - wf_output = Output(lp_task.outputs.out, sdk_type=_Types.Integer) diff --git a/tests/flytekit/common/workflows/failing_workflows.py b/tests/flytekit/common/workflows/failing_workflows.py deleted file mode 100644 index 260e460b17..0000000000 --- a/tests/flytekit/common/workflows/failing_workflows.py +++ /dev/null @@ -1,28 +0,0 @@ -from flytekit.models.core.workflow import WorkflowMetadata -from flytekit.sdk.tasks import python_task -from flytekit.sdk.workflow import workflow_class - - -@python_task -def div_zero(wf_params): - return 5 / 0 - - -@python_task -def log_something(wf_params): - wf_params.logging.warn("Hello world") - - -@workflow_class(on_failure=WorkflowMetadata.OnFailurePolicy.FAIL_AFTER_EXECUTABLE_NODES_COMPLETE) -class FailingWorkflowWithRunToCompletion(object): - """ - [start] -> [first_layer] -> [second_layer] -> [end] - \\_ [first_layer_2] _/ - """ - - first_layer = log_something() - first_layer_2 = div_zero() - second_layer = div_zero() - - # This forces second_layer node to run after first layer - first_layer >> second_layer diff --git a/tests/flytekit/common/workflows/gpu.py b/tests/flytekit/common/workflows/gpu.py deleted file mode 100644 index 7e7621ac69..0000000000 --- a/tests/flytekit/common/workflows/gpu.py +++ /dev/null @@ -1,20 +0,0 @@ -from flytekit.sdk.tasks import inputs, outputs, python_task -from flytekit.sdk.types import Types -from flytekit.sdk.workflow import Input, Output, workflow_class - - -@inputs(a=Types.Integer) -@outputs(b=Types.Integer) -@python_task(gpu_request="1", gpu_limit="1") -def add_one(wf_params, a, b): - # TODO lets add a test that works with tensorflow, but we need it to be in - # a different container - b.set(a + 1) - - -@workflow_class -class SimpleWorkflow(object): - input_1 = Input(Types.Integer) - input_2 = Input(Types.Integer, default=5, help="Not required.") - a = add_one(a=input_1) - output = Output(a.outputs.b, sdk_type=Types.Integer) diff --git a/tests/flytekit/common/workflows/hive.py b/tests/flytekit/common/workflows/hive.py deleted file mode 100644 index b3cb56fe5d..0000000000 --- a/tests/flytekit/common/workflows/hive.py +++ /dev/null @@ -1,33 +0,0 @@ -import six as _six - -from flytekit.sdk.tasks import inputs, outputs, python_task, qubole_hive_task -from flytekit.sdk.types import Types -from flytekit.sdk.workflow import workflow_class - - -@outputs(hive_results=[Types.Schema()]) -@qubole_hive_task(tags=[_six.text_type("these"), _six.text_type("are"), _six.text_type("tags")]) -def generate_queries(wf_params, hive_results): - q1 = "SELECT 1" - q2 = "SELECT 'two'" - schema_1, formatted_query_1 = Types.Schema().create_from_hive_query(select_query=q1) - schema_2, formatted_query_2 = Types.Schema().create_from_hive_query(select_query=q2) - - hive_results.set([schema_1, schema_2]) - return [formatted_query_1, formatted_query_2] - - -@inputs(ss=[Types.Schema()]) -@python_task -def print_schemas(wf_params, ss): - for s in ss: - with s as r: - for df in r.iter_chunks(): - df = r.read() - print(df) - - -@workflow_class -class ExampleQueryWorkflow(object): - a = generate_queries() - b = print_schemas(ss=a.outputs.hive_results) diff --git a/tests/flytekit/common/workflows/nested.py b/tests/flytekit/common/workflows/nested.py deleted file mode 100644 index a364356ff6..0000000000 --- a/tests/flytekit/common/workflows/nested.py +++ /dev/null @@ -1,49 +0,0 @@ -from flytekit.sdk.tasks import inputs, outputs, python_task -from flytekit.sdk.types import Types -from flytekit.sdk.workflow import Input, Output, workflow_class - - -@inputs(a=Types.Integer) -@outputs(b=Types.Integer) -@python_task -def add_one(wf_params, a, b): - b.set(a + 1) - - -@inputs(a=Types.Integer) -@outputs(b=Types.Integer) -@python_task -def subtract_one(wf_params, a, b): - b.set(a - 1) - - -@inputs(a=Types.Integer, b=Types.Integer) -@outputs(c=Types.Integer) -@python_task -def sum(wf_params, a, b, c): - c.set(a + b) - - -@workflow_class -class Child(object): - input_1 = Input(Types.Integer) - input_2 = Input(Types.Integer, default=5, help="Not required.") - a = add_one(a=input_1) - b = add_one(a=input_2) - c = add_one(a=100) - output = Output(c.outputs.b, sdk_type=Types.Integer) - - -# Create a simple launch plan without any overrides for inputs to the child workflow. -child_lp = Child.create_launch_plan() - - -# Create a parent workflow that invokes the child workflow previously declared. Note that it takes advantage of -# default and required inputs on different calls. -@workflow_class -class Parent(object): - input_1 = Input(Types.Integer) - child1 = child_lp(input_1=input_1) - child2 = child_lp(input_1=input_1, input_2=10) - final_sum = sum(a=child1.outputs.output, b=child2.outputs.output) - output = Output(final_sum.outputs.c, sdk_type=Types.Integer) diff --git a/tests/flytekit/common/workflows/notebook.py b/tests/flytekit/common/workflows/notebook.py deleted file mode 100644 index e8af3b7467..0000000000 --- a/tests/flytekit/common/workflows/notebook.py +++ /dev/null @@ -1,25 +0,0 @@ -from flytekit.contrib.notebook.tasks import python_notebook, spark_notebook -from flytekit.sdk.tasks import inputs, outputs -from flytekit.sdk.types import Types -from flytekit.sdk.workflow import Input, workflow_class - -interactive_python = python_notebook( - notebook_path="../../../../notebook-task-examples/python-notebook.ipynb", - inputs=inputs(pi=Types.Float), - outputs=outputs(out=Types.Float), - cpu_request="1", - memory_request="1G", -) - -interactive_spark = spark_notebook( - notebook_path="../../../../notebook-task-examples/spark-notebook-pi.ipynb", - inputs=inputs(partitions=Types.Integer), - outputs=outputs(pi=Types.Float), -) - - -@workflow_class -class FlyteNotebookSparkWorkflow(object): - partitions = Input(Types.Integer, default=10) - out1 = interactive_spark(partitions=partitions) - out2 = interactive_python(pi=out1.outputs.pi) diff --git a/tests/flytekit/common/workflows/notifications.py b/tests/flytekit/common/workflows/notifications.py deleted file mode 100644 index d09bc14fad..0000000000 --- a/tests/flytekit/common/workflows/notifications.py +++ /dev/null @@ -1,34 +0,0 @@ -from flytekit.common import notifications as _notifications -from flytekit.models.core import execution as _execution -from flytekit.sdk.tasks import inputs, outputs, python_task -from flytekit.sdk.types import Types -from flytekit.sdk.workflow import Input, workflow_class - - -@inputs(a=Types.Integer, b=Types.Integer) -@outputs(c=Types.Integer) -@python_task -def add_two_integers(wf_params, a, b, c): - c.set(a + b) - - -@workflow_class -class BasicWorkflow(object): - input_1 = Input(Types.Integer) - input_2 = Input(Types.Integer, default=1, help="Not required.") - a = add_two_integers(a=input_1, b=input_2) - - -notification_lp = BasicWorkflow.create_launch_plan( - notifications=[ - _notifications.Email( - [ - _execution.WorkflowExecutionPhase.SUCCEEDED, - _execution.WorkflowExecutionPhase.FAILED, - _execution.WorkflowExecutionPhase.TIMED_OUT, - _execution.WorkflowExecutionPhase.ABORTED, - ], - ["flyte-test-notifications@mydomain.com"], - ) - ] -) diff --git a/tests/flytekit/common/workflows/presto.py b/tests/flytekit/common/workflows/presto.py deleted file mode 100644 index e681a3e953..0000000000 --- a/tests/flytekit/common/workflows/presto.py +++ /dev/null @@ -1,25 +0,0 @@ -from flytekit.common.tasks.presto_task import SdkPrestoTask -from flytekit.sdk.tasks import inputs -from flytekit.sdk.types import Types -from flytekit.sdk.workflow import Input, Output, workflow_class - -schema = Types.Schema([("a", Types.String), ("b", Types.Integer)]) - -presto_task = SdkPrestoTask( - task_inputs=inputs(ds=Types.String, rg=Types.String), - statement="SELECT * FROM hive.city.fact_airport_sessions WHERE ds = '{{ .Inputs.ds}}' LIMIT 10", - output_schema=schema, - routing_group="{{ .Inputs.rg }}", - # catalog="hive", - # schema="city", -) - - -@workflow_class() -class PrestoWorkflow(object): - ds = Input(Types.String, required=True, help="Test string with no default") - # routing_group = Input(Types.String, required=True, help="Test string with no default") - - p_task = presto_task(ds=ds, rg="etl") - - output_a = Output(p_task.outputs.results, sdk_type=schema) diff --git a/tests/flytekit/common/workflows/python.py b/tests/flytekit/common/workflows/python.py deleted file mode 100644 index a0d423b86c..0000000000 --- a/tests/flytekit/common/workflows/python.py +++ /dev/null @@ -1,67 +0,0 @@ -from flytekit.sdk.tasks import inputs, outputs, python_task -from flytekit.sdk.types import Types -from flytekit.sdk.workflow import Input, workflow_class - - -@inputs(value_to_print=Types.Integer) -@outputs(out=Types.Integer) -@python_task(cache_version="1") -def add_one_and_print(workflow_parameters, value_to_print, out): - workflow_parameters.stats.incr("task_run") - added = value_to_print + 1 - print("My printed value: {}".format(added)) - out.set(added) - - -@inputs(value1_to_print=Types.Integer, value2_to_print=Types.Integer) -@outputs(out=Types.Integer) -@python_task(cache_version="1") -def sum_non_none(workflow_parameters, value1_to_print, value2_to_print, out): - workflow_parameters.stats.incr("task_run") - added = 0 - for value in [value1_to_print, value2_to_print]: - print("Adding values: {}".format(value)) - if value is not None: - added += value - added += 1 - print("My printed value: {}".format(added)) - out.set(added) - - -@inputs( - value1_to_add=Types.Integer, - value2_to_add=Types.Integer, - value3_to_add=Types.Integer, - value4_to_add=Types.Integer, -) -@outputs(out=Types.Integer) -@python_task(cache_version="1") -def sum_and_print(workflow_parameters, value1_to_add, value2_to_add, value3_to_add, value4_to_add, out): - workflow_parameters.stats.incr("task_run") - summed = sum([value1_to_add, value2_to_add, value3_to_add, value4_to_add]) - print("Summed up to: {}".format(summed)) - out.set(summed) - - -@inputs(value_to_print=Types.Integer, date_triggered=Types.Datetime) -@python_task(cache_version="1") -def print_every_time(workflow_parameters, value_to_print, date_triggered): - workflow_parameters.stats.incr("task_run") - print("My printed value: {} @ {}".format(value_to_print, date_triggered)) - - -@workflow_class -class PythonTasksWorkflow(object): - triggered_date = Input(Types.Datetime) - print1a = add_one_and_print(value_to_print=3) - print1b = add_one_and_print(value_to_print=101) - print2 = sum_non_none(value1_to_print=print1a.outputs.out, value2_to_print=print1b.outputs.out) - print3 = add_one_and_print(value_to_print=print2.outputs.out) - print4 = add_one_and_print(value_to_print=print3.outputs.out) - print_sum = sum_and_print( - value1_to_add=print2.outputs.out, - value2_to_add=print3.outputs.out, - value3_to_add=print4.outputs.out, - value4_to_add=100, - ) - print_always = print_every_time(value_to_print=print_sum.outputs.out, date_triggered=triggered_date) diff --git a/tests/flytekit/common/workflows/raw_container.py b/tests/flytekit/common/workflows/raw_container.py deleted file mode 100644 index 60edbcc737..0000000000 --- a/tests/flytekit/common/workflows/raw_container.py +++ /dev/null @@ -1,31 +0,0 @@ -from flytekit.common.tasks.raw_container import SdkRawContainerTask -from flytekit.sdk.types import Types -from flytekit.sdk.workflow import Input, Output, workflow_class - -square = SdkRawContainerTask( - input_data_dir="/var/inputs", - output_data_dir="/var/outputs", - inputs={"val": Types.Integer}, - outputs={"out": Types.Integer}, - image="alpine", - command=["sh", "-c", "echo $(( {{.Inputs.val}} * {{.Inputs.val}} )) | tee /var/outputs/out"], -) - -sum = SdkRawContainerTask( - input_data_dir="/var/flyte/inputs", - output_data_dir="/var/flyte/outputs", - inputs={"x": Types.Integer, "y": Types.Integer}, - outputs={"out": Types.Integer}, - image="alpine", - command=["sh", "-c", "echo $(( {{.Inputs.x}} + {{.Inputs.y}} )) | tee /var/flyte/outputs/out"], -) - - -@workflow_class -class RawContainerWorkflow(object): - val1 = Input(Types.Integer) - val2 = Input(Types.Integer) - sq1 = square(val=val1) - sq2 = square(val=val2) - sm = sum(x=sq1.outputs.out, y=sq2.outputs.out) - sum_of_squares = Output(sm.outputs.out, sdk_type=Types.Integer) diff --git a/tests/flytekit/common/workflows/raw_edge_detector.py b/tests/flytekit/common/workflows/raw_edge_detector.py deleted file mode 100644 index eeefac4de1..0000000000 --- a/tests/flytekit/common/workflows/raw_edge_detector.py +++ /dev/null @@ -1,20 +0,0 @@ -from flytekit.common.tasks.raw_container import SdkRawContainerTask -from flytekit.sdk.types import Types -from flytekit.sdk.workflow import Input, Output, workflow_class - -edges = SdkRawContainerTask( - input_data_dir="/inputs", - output_data_dir="/outputs", - inputs={"image": Types.Blob, "script": Types.Blob}, - outputs={"edges": Types.Blob}, - image="jjanzic/docker-python3-opencv", - command=["python", "{{.inputs.script}}", "/inputs/image", "/outputs/edges"], -) - - -@workflow_class -class EdgeDetector(object): - script = Input(Types.Blob) - image = Input(Types.Blob) - edge_task = edges(script=script, image=image) - out = Output(edge_task.outputs.edges, sdk_type=Types.Blob) diff --git a/tests/flytekit/common/workflows/scala_spark.py b/tests/flytekit/common/workflows/scala_spark.py deleted file mode 100644 index 7df3b0fb32..0000000000 --- a/tests/flytekit/common/workflows/scala_spark.py +++ /dev/null @@ -1,32 +0,0 @@ -from flytekit.sdk.spark_types import SparkType -from flytekit.sdk.tasks import generic_spark_task, inputs, python_task -from flytekit.sdk.types import Types -from flytekit.sdk.workflow import Input, workflow_class - -scala_spark = generic_spark_task( - spark_type=SparkType.SCALA, - inputs=inputs(partitions=Types.Integer), - main_class="org.apache.spark.examples.SparkPi", - main_application_file="local:///opt/spark/examples/jars/spark-examples.jar", - spark_conf={ - "spark.driver.memory": "1000M", - "spark.executor.memory": "1000M", - "spark.executor.cores": "1", - "spark.executor.instances": "2", - }, - cache_version="1", -) - - -@inputs(date_triggered=Types.Datetime) -@python_task(cache_version="1") -def print_every_time(workflow_parameters, date_triggered): - print("My input : {}".format(date_triggered)) - - -@workflow_class -class SparkTasksWorkflow(object): - triggered_date = Input(Types.Datetime) - partitions = Input(Types.Integer) - spark_task = scala_spark(partitions=partitions) - print_always = print_every_time(date_triggered=triggered_date) diff --git a/tests/flytekit/common/workflows/sidecar.py b/tests/flytekit/common/workflows/sidecar.py deleted file mode 100644 index ff0d4e7484..0000000000 --- a/tests/flytekit/common/workflows/sidecar.py +++ /dev/null @@ -1,56 +0,0 @@ -import os -import time - -from k8s.io.api.core.v1 import generated_pb2 - -from flytekit.sdk.tasks import sidecar_task -from flytekit.sdk.types import Types -from flytekit.sdk.workflow import Input, workflow_class - - -def generate_pod_spec_for_task(): - pod_spec = generated_pb2.PodSpec() - secondary_container = generated_pb2.Container( - name="secondary", - image="alpine", - ) - secondary_container.command.extend(["/bin/sh"]) - secondary_container.args.extend(["-c", "echo hi sidecar world > /data/message.txt"]) - shared_volume_mount = generated_pb2.VolumeMount( - name="shared-data", - mountPath="/data", - ) - secondary_container.volumeMounts.extend([shared_volume_mount]) - - primary_container = generated_pb2.Container(name="primary") - primary_container.volumeMounts.extend([shared_volume_mount]) - - pod_spec.volumes.extend( - [ - generated_pb2.Volume( - name="shared-data", - volumeSource=generated_pb2.VolumeSource( - emptyDir=generated_pb2.EmptyDirVolumeSource( - medium="Memory", - ) - ), - ) - ] - ) - pod_spec.containers.extend([primary_container, secondary_container]) - return pod_spec - - -@sidecar_task( - pod_spec=generate_pod_spec_for_task(), - primary_container_name="primary", -) -def a_sidecar_task(wfparams): - while not os.path.isfile("/data/message.txt"): - time.sleep(5) - - -@workflow_class -class SimpleSidecarWorkflow(object): - input_1 = Input(Types.String) - my_sidecar_task = a_sidecar_task() diff --git a/tests/flytekit/common/workflows/simple.py b/tests/flytekit/common/workflows/simple.py deleted file mode 100644 index f264fe39bf..0000000000 --- a/tests/flytekit/common/workflows/simple.py +++ /dev/null @@ -1,114 +0,0 @@ -import pandas as _pd - -from flytekit.sdk.tasks import inputs, outputs, python_task -from flytekit.sdk.types import Types -from flytekit.sdk.workflow import Input, workflow_class - - -@inputs(a=Types.Integer) -@outputs(b=Types.Integer) -@python_task -def add_one(wf_params, a, b): - b.set(a + 1) - - -@inputs(a=Types.Integer) -@outputs(b=Types.Integer) -@python_task(cache=True, cache_version="1") -def subtract_one(wf_params, a, b): - b.set(a - 1) - - -@outputs( - a=Types.Blob, - b=Types.CSV, - c=Types.MultiPartCSV, - d=Types.MultiPartBlob, - e=Types.Schema([("a", Types.Integer), ("b", Types.Integer)]), -) -@python_task -def write_special_types(wf_params, a, b, c, d, e): - blob = Types.Blob() - with blob as w: - w.write("hello I'm a blob".encode("utf-8")) - - csv = Types.CSV() - with csv as w: - w.write("hello,i,iz,blob") - - mpcsv = Types.MultiPartCSV() - with mpcsv.create_part("000000") as w: - w.write("hello,i,iz,blob") - with mpcsv.create_part("000001") as w: - w.write("hello,i,iz,blob2") - - mpblob = Types.MultiPartBlob() - with mpblob.create_part("000000") as w: - w.write("hello I'm a mp blob".encode("utf-8")) - with mpblob.create_part("000001") as w: - w.write("hello I'm a mp blob too".encode("utf-8")) - - schema = Types.Schema([("a", Types.Integer), ("b", Types.Integer)])() - with schema as w: - w.write(_pd.DataFrame.from_dict({"a": [1, 2, 3], "b": [4, 5, 6]})) - w.write(_pd.DataFrame.from_dict({"a": [3, 2, 1], "b": [6, 5, 4]})) - - a.set(blob) - b.set(csv) - c.set(mpcsv) - d.set(mpblob) - e.set(schema) - - -@inputs( - a=Types.Blob, - b=Types.CSV, - c=Types.MultiPartCSV, - d=Types.MultiPartBlob, - e=Types.Schema([("a", Types.Integer), ("b", Types.Integer)]), -) -@python_task -def read_special_types(wf_params, a, b, c, d, e): - with a as r: - assert r.read().decode("utf-8") == "hello I'm a blob" - - with b as r: - assert r.read() == "hello,i,iz,blob" - - with c as r: - assert len(r) == 2 - assert r[0].read() == "hello,i,iz,blob" - assert r[1].read() == "hello,i,iz,blob2" - - with d as r: - assert len(r) == 2 - assert r[0].read().decode("utf-8") == "hello I'm a mp blob" - assert r[1].read().decode("utf-8") == "hello I'm a mp blob too" - - with e as r: - df = r.read() - assert df["a"].tolist() == [1, 2, 3] - assert df["b"].tolist() == [4, 5, 6] - - df = r.read() - assert df["a"].tolist() == [3, 2, 1] - assert df["b"].tolist() == [6, 5, 4] - assert r.read() is None - - -@workflow_class -class SimpleWorkflow(object): - input_1 = Input(Types.Integer) - input_2 = Input(Types.Integer, default=5, help="Not required.") - a = add_one(a=input_1) - b = add_one(a=input_2) - c = subtract_one(a=input_1) - - d = write_special_types() - e = read_special_types( - a=d.outputs.a, - b=d.outputs.b, - c=d.outputs.c, - d=d.outputs.d, - e=d.outputs.e, - ) diff --git a/tests/flytekit/common/workflows/spark.py b/tests/flytekit/common/workflows/spark.py deleted file mode 100644 index b3f381f6ba..0000000000 --- a/tests/flytekit/common/workflows/spark.py +++ /dev/null @@ -1,50 +0,0 @@ -import random -from operator import add - -from six.moves import range - -from flytekit.sdk.tasks import inputs, outputs, python_task, spark_task -from flytekit.sdk.types import Types -from flytekit.sdk.workflow import Input, workflow_class - - -@inputs(partitions=Types.Integer) -@outputs(out=Types.Float) -@spark_task( - spark_conf={ - "spark.driver.memory": "1000M", - "spark.executor.memory": "1000M", - "spark.executor.cores": "1", - "spark.executor.instances": "2", - "spark.hadoop.mapred.output.committer.class": "org.apache.hadoop.mapred.DirectFileOutputCommitter", - "spark.hadoop.mapreduce.use.directfileoutputcommitter": "true", - }, - cache_version="1", -) -def hello_spark(workflow_parameters, spark_context, partitions, out): - print("Starting Spark with Partitions: {}".format(partitions)) - - n = 100000 * partitions - count = spark_context.parallelize(range(1, n + 1), partitions).map(f).reduce(add) - pi_val = 4.0 * count / n - print("Pi val is :{}".format(pi_val)) - out.set(pi_val) - - -@inputs(value_to_print=Types.Float, date_triggered=Types.Datetime) -@python_task(cache_version="1") -def print_every_time(workflow_parameters, value_to_print, date_triggered): - print("My printed value: {} @ {}".format(value_to_print, date_triggered)) - - -def f(_): - x = random.random() * 2 - 1 - y = random.random() * 2 - 1 - return 1 if x ** 2 + y ** 2 <= 1 else 0 - - -@workflow_class -class SparkTasksWorkflow(object): - triggered_date = Input(Types.Datetime) - sparkTask = hello_spark(partitions=50) - print_always = print_every_time(value_to_print=sparkTask.outputs.out, date_triggered=triggered_date) diff --git a/tests/flytekit/integration/remote/test_remote.py b/tests/flytekit/integration/remote/test_remote.py index a45c037279..7784d76d82 100644 --- a/tests/flytekit/integration/remote/test_remote.py +++ b/tests/flytekit/integration/remote/test_remote.py @@ -9,8 +9,8 @@ import pytest from flytekit import kwtypes -from flytekit.common.exceptions.user import FlyteAssertion, FlyteEntityNotExistException from flytekit.core.launch_plan import LaunchPlan +from flytekit.exceptions.user import FlyteAssertion, FlyteEntityNotExistException from flytekit.extras.sqlite3.task import SQLite3Config, SQLite3Task from flytekit.remote.remote import FlyteRemote from flytekit.types.schema import FlyteSchema diff --git a/tests/flytekit/loadtests/__init__.py b/tests/flytekit/loadtests/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/tests/flytekit/loadtests/cp_orchestrator.py b/tests/flytekit/loadtests/cp_orchestrator.py deleted file mode 100644 index 776afecd59..0000000000 --- a/tests/flytekit/loadtests/cp_orchestrator.py +++ /dev/null @@ -1,29 +0,0 @@ -from six.moves import range - -from flytekit.sdk.workflow import workflow_class -from tests.flytekit.loadtests.cp_python import FlyteCPPythonLoadTestWorkflow -from tests.flytekit.loadtests.cp_spark import FlyteCPSparkLoadTestWorkflow - -# launch plans for individual load tests. -python_loadtest_lp = FlyteCPPythonLoadTestWorkflow.create_launch_plan() -spark_loadtest_lp = FlyteCPSparkLoadTestWorkflow.create_launch_plan() - - -# Orchestrator workflow invokes the individual load test workflows (hive, python, spark). Its static for now but we -# will make it dynamic in future -@workflow_class -class CPLoadTestOrchestrationWorkflow(object): - - # python load tests. 5 tasks each. Total: 1 cpu 5gb memory per python workflow. - python_task_count = 50 - p = [None] * python_task_count - - for i in range(0, python_task_count): - p[i] = python_loadtest_lp() - - # spark load tests. - spark_task_count = 30 - s = [None] * spark_task_count - - for i in range(0, spark_task_count): - s[i] = spark_loadtest_lp() diff --git a/tests/flytekit/loadtests/cp_python.py b/tests/flytekit/loadtests/cp_python.py deleted file mode 100644 index f26cbb51f7..0000000000 --- a/tests/flytekit/loadtests/cp_python.py +++ /dev/null @@ -1,28 +0,0 @@ -import time - -from six.moves import range - -from flytekit.sdk.tasks import inputs, outputs, python_task -from flytekit.sdk.types import Types -from flytekit.sdk.workflow import workflow_class - - -@inputs(value1_to_add=Types.Integer, value2_to_add=Types.Integer) -@outputs(out=Types.Integer) -@python_task(cpu_request="1", cpu_limit="1", memory_request="3G") -def sum_and_print(workflow_parameters, value1_to_add, value2_to_add, out): - for i in range(5 * 60): - print("This is load test task. I have been running for {} seconds.".format(i)) - time.sleep(1) - - summed = sum([value1_to_add, value2_to_add]) - print("Summed up to: {}".format(summed)) - out.set(summed) - - -@workflow_class -class FlyteCPPythonLoadTestWorkflow(object): - - print_sum = [None] * 5 - for i in range(0, 5): - print_sum[i] = sum_and_print(value1_to_add=1, value2_to_add=1) diff --git a/tests/flytekit/loadtests/cp_spark.py b/tests/flytekit/loadtests/cp_spark.py deleted file mode 100644 index b79fbc3509..0000000000 --- a/tests/flytekit/loadtests/cp_spark.py +++ /dev/null @@ -1,42 +0,0 @@ -import random -from operator import add - -from six.moves import range - -from flytekit.sdk.tasks import inputs, outputs, spark_task -from flytekit.sdk.types import Types -from flytekit.sdk.workflow import workflow_class - - -@inputs(partitions=Types.Integer) -@outputs(out=Types.Float) -@spark_task( - spark_conf={ - "spark.driver.memory": "600M", - "spark.executor.memory": "600M", - "spark.executor.cores": "1", - "spark.executor.instances": "1", - "spark.hadoop.mapred.output.committer.class": "org.apache.hadoop.mapred.DirectFileOutputCommitter", - "spark.hadoop.mapreduce.use.directfileoutputcommitter": "true", - }, - cache_version="1", -) -def hello_spark(workflow_parameters, spark_context, partitions, out): - print("Starting Spark with Partitions: {}".format(partitions)) - - n = 30000 * partitions - count = spark_context.parallelize(range(1, n + 1), partitions).map(f).reduce(add) - pi_val = 4.0 * count / n - print("Pi val is :{}".format(pi_val)) - out.set(pi_val) - - -def f(_): - x = random.random() * 2 - 1 - y = random.random() * 2 - 1 - return 1 if x ** 2 + y ** 2 <= 1 else 0 - - -@workflow_class -class FlyteCPSparkLoadTestWorkflow(object): - sparkTask = hello_spark(partitions=50) diff --git a/tests/flytekit/loadtests/dynamic_job.py b/tests/flytekit/loadtests/dynamic_job.py deleted file mode 100644 index 520ff401a7..0000000000 --- a/tests/flytekit/loadtests/dynamic_job.py +++ /dev/null @@ -1,40 +0,0 @@ -import time - -from six.moves import range - -from flytekit.sdk.tasks import dynamic_task, inputs, outputs, python_task -from flytekit.sdk.types import Types -from flytekit.sdk.workflow import Input, workflow_class - - -@inputs(value1=Types.Integer) -@outputs(out=Types.Integer) -@python_task(cpu_request="1", cpu_limit="1", memory_request="5G") -def dynamic_sub_task(workflow_parameters, value1, out): - for i in range(11 * 60): - print("This is load test task. I have been running for {} seconds.".format(i)) - time.sleep(1) - - output = value1 * 2 - print("Output: {}".format(output)) - out.set(output) - - -@inputs(tasks_count=Types.Integer) -@outputs(out=[Types.Integer]) -@dynamic_task(cache_version="1") -def dynamic_task(workflow_parameters, tasks_count, out): - res = [] - for i in range(0, tasks_count): - task = dynamic_sub_task(value1=i) - yield task - res.append(task.outputs.out) - - # Define how to set the final result of the task - out.set(res) - - -@workflow_class -class FlyteDJOLoadTestWorkflow(object): - tasks_count = Input(Types.Integer) - dj = dynamic_task(tasks_count=tasks_count) diff --git a/tests/flytekit/loadtests/orchestrator.py b/tests/flytekit/loadtests/orchestrator.py deleted file mode 100644 index 5981502319..0000000000 --- a/tests/flytekit/loadtests/orchestrator.py +++ /dev/null @@ -1,48 +0,0 @@ -from six.moves import range - -from flytekit.sdk.types import Types -from flytekit.sdk.workflow import Input, workflow_class -from tests.flytekit.loadtests.dynamic_job import FlyteDJOLoadTestWorkflow -from tests.flytekit.loadtests.hive import FlyteHiveLoadTestWorkflow -from tests.flytekit.loadtests.python import FlytePythonLoadTestWorkflow -from tests.flytekit.loadtests.spark import FlyteSparkLoadTestWorkflow - -# launch plans for individual load tests. -python_loadtest_lp = FlytePythonLoadTestWorkflow.create_launch_plan() -hive_loadtest_lp = FlyteHiveLoadTestWorkflow.create_launch_plan() -spark_loadtest_lp = FlyteSparkLoadTestWorkflow.create_launch_plan() -dynamic_job_loadtest_lp = FlyteDJOLoadTestWorkflow.create_launch_plan() - - -# Orchestrator workflow invokes the individual load test workflows (hive, python, spark). Its static for now but we -# will make it dynamic in future,. -@workflow_class -class LoadTestOrchestrationWorkflow(object): - - # 30 python tasks ~= 75 i3.16x nodes on AWS Batch - python_task_count = 30 - # 30 spark tasks ~= 60 i3.16x nodes on AWS Batch - spark_task_count = 30 - # 3 dynamic-jobs each of 1000 tasks ~= 3*20 i3.16x nodes on AWS Batch - djo_task_count = 1000 - dj_count = 3 - - p = [None] * python_task_count - s = [None] * spark_task_count - d = [None] * dj_count - - # python tasks - for i in range(0, python_task_count): - p[i] = python_loadtest_lp() - - # dynamic-job tasks - for i in range(0, dj_count): - d[i] = dynamic_job_loadtest_lp(tasks_count=djo_task_count) - - # hive load tests. - # h1 = hive_loadtest_lp() - - # spark load tests - trigger_time = Input(Types.Datetime) - for i in range(0, spark_task_count): - s[i] = spark_loadtest_lp(triggered_date=trigger_time, offset=i) diff --git a/tests/flytekit/loadtests/python.py b/tests/flytekit/loadtests/python.py deleted file mode 100644 index e6da8df722..0000000000 --- a/tests/flytekit/loadtests/python.py +++ /dev/null @@ -1,27 +0,0 @@ -import time - -from six.moves import range - -from flytekit.sdk.tasks import inputs, outputs, python_task -from flytekit.sdk.types import Types -from flytekit.sdk.workflow import workflow_class - - -@inputs(value1_to_add=Types.Integer, value2_to_add=Types.Integer) -@outputs(out=Types.Integer) -@python_task(cpu_request="5", cpu_limit="5", memory_request="32G") -def sum_and_print(workflow_parameters, value1_to_add, value2_to_add, out): - for i in range(11 * 60): - print("This is load test task. I have been running for {} seconds.".format(i)) - time.sleep(1) - - summed = sum([value1_to_add, value2_to_add]) - print("Summed up to: {}".format(summed)) - out.set(summed) - - -@workflow_class -class FlytePythonLoadTestWorkflow(object): - print_sum = [None] * 30 - for i in range(0, 30): - print_sum[i] = sum_and_print(value1_to_add=1, value2_to_add=1) diff --git a/tests/flytekit/unit/bin/test_python_entrypoint.py b/tests/flytekit/unit/bin/test_python_entrypoint.py index d54d7dc52b..3a71b567d8 100644 --- a/tests/flytekit/unit/bin/test_python_entrypoint.py +++ b/tests/flytekit/unit/bin/test_python_entrypoint.py @@ -1,194 +1,30 @@ -import os import typing from collections import OrderedDict import mock import pytest -import six -from click.testing import CliRunner -from flyteidl.core import literals_pb2 as _literals_pb2 from flyteidl.core.errors_pb2 import ErrorDocument -from flytekit.bin.entrypoint import _dispatch_execute, _legacy_execute_task, execute_task_cmd, setup_execution -from flytekit.common import constants as _constants -from flytekit.common import utils as _utils -from flytekit.common.exceptions import user as user_exceptions -from flytekit.common.exceptions.scopes import system_entry_point -from flytekit.common.types import helpers as _type_helpers -from flytekit.configuration import TemporaryConfiguration as _TemporaryConfiguration +from flytekit.bin.entrypoint import _dispatch_execute, setup_execution from flytekit.core import context_manager from flytekit.core.base_task import IgnoreOutputs from flytekit.core.dynamic_workflow_task import dynamic from flytekit.core.promise import VoidPromise from flytekit.core.task import task from flytekit.core.type_engine import TypeEngine +from flytekit.exceptions import user as user_exceptions +from flytekit.exceptions.scopes import system_entry_point from flytekit.extras.persistence.gcs_gsutil import GCSPersistence from flytekit.extras.persistence.s3_awscli import S3Persistence from flytekit.models import literals as _literal_models from flytekit.models.core import errors as error_models from flytekit.models.core import execution as execution_models -from tests.flytekit.common import task_definitions as _task_defs - - -def _type_map_from_variable_map(variable_map): - return {k: _type_helpers.get_sdk_type_from_literal_type(v.type) for k, v in six.iteritems(variable_map)} - - -def test_single_step_entrypoint_in_proc(): - with _TemporaryConfiguration( - os.path.join(os.path.dirname(__file__), "fake.config"), - internal_overrides={"project": "test", "domain": "development"}, - ): - with _utils.AutoDeletingTempDir("in") as input_dir: - literal_map = _type_helpers.pack_python_std_map_to_literal_map( - {"a": 9}, - _type_map_from_variable_map(_task_defs.add_one.interface.inputs), - ) - input_file = os.path.join(input_dir.name, "inputs.pb") - _utils.write_proto_to_file(literal_map.to_flyte_idl(), input_file) - - with _utils.AutoDeletingTempDir("out") as output_dir: - _legacy_execute_task( - _task_defs.add_one.task_module, - _task_defs.add_one.task_function_name, - input_file, - output_dir.name, - output_dir.name, - False, - ) - - p = _utils.load_proto_from_file( - _literals_pb2.LiteralMap, - os.path.join(output_dir.name, _constants.OUTPUT_FILE_NAME), - ) - raw_map = _type_helpers.unpack_literal_map_to_sdk_python_std( - _literal_models.LiteralMap.from_flyte_idl(p), - _type_map_from_variable_map(_task_defs.add_one.interface.outputs), - ) - assert raw_map["b"] == 10 - assert len(raw_map) == 1 - - -def test_single_step_entrypoint_out_of_proc(): - with _TemporaryConfiguration( - os.path.join(os.path.dirname(__file__), "fake.config"), - internal_overrides={"project": "test", "domain": "development"}, - ): - with _utils.AutoDeletingTempDir("in") as input_dir: - literal_map = _type_helpers.pack_python_std_map_to_literal_map( - {"a": 9}, - _type_map_from_variable_map(_task_defs.add_one.interface.inputs), - ) - input_file = os.path.join(input_dir.name, "inputs.pb") - _utils.write_proto_to_file(literal_map.to_flyte_idl(), input_file) - - with _utils.AutoDeletingTempDir("out") as output_dir: - cmd = [] - cmd.extend(["--task-module", _task_defs.add_one.task_module]) - cmd.extend(["--task-name", _task_defs.add_one.task_function_name]) - cmd.extend(["--inputs", input_file]) - cmd.extend(["--output-prefix", output_dir.name]) - result = CliRunner().invoke(execute_task_cmd, cmd) - - assert result.exit_code == 0 - p = _utils.load_proto_from_file( - _literals_pb2.LiteralMap, - os.path.join(output_dir.name, _constants.OUTPUT_FILE_NAME), - ) - raw_map = _type_helpers.unpack_literal_map_to_sdk_python_std( - _literal_models.LiteralMap.from_flyte_idl(p), - _type_map_from_variable_map(_task_defs.add_one.interface.outputs), - ) - assert raw_map["b"] == 10 - assert len(raw_map) == 1 - - -def test_arrayjob_entrypoint_in_proc(): - with _TemporaryConfiguration( - os.path.join(os.path.dirname(__file__), "fake.config"), - internal_overrides={"project": "test", "domain": "development"}, - ): - with _utils.AutoDeletingTempDir("dir") as dir: - literal_map = _type_helpers.pack_python_std_map_to_literal_map( - {"a": 9}, - _type_map_from_variable_map(_task_defs.add_one.interface.inputs), - ) - - input_dir = os.path.join(dir.name, "1") - os.mkdir(input_dir) # auto cleanup will take this subdir into account - - input_file = os.path.join(input_dir, "inputs.pb") - _utils.write_proto_to_file(literal_map.to_flyte_idl(), input_file) - - # construct indexlookup.pb which has array: [1] - mapped_index = _literal_models.Literal( - _literal_models.Scalar(primitive=_literal_models.Primitive(integer=1)) - ) - index_lookup_collection = _literal_models.LiteralCollection([mapped_index]) - index_lookup_file = os.path.join(dir.name, "indexlookup.pb") - _utils.write_proto_to_file(index_lookup_collection.to_flyte_idl(), index_lookup_file) - - # fake arrayjob task by setting environment variables - orig_env_index_var_name = os.environ.get("BATCH_JOB_ARRAY_INDEX_VAR_NAME") - orig_env_array_index = os.environ.get("AWS_BATCH_JOB_ARRAY_INDEX") - os.environ["BATCH_JOB_ARRAY_INDEX_VAR_NAME"] = "AWS_BATCH_JOB_ARRAY_INDEX" - os.environ["AWS_BATCH_JOB_ARRAY_INDEX"] = "0" - - _legacy_execute_task( - _task_defs.add_one.task_module, - _task_defs.add_one.task_function_name, - dir.name, - dir.name, - dir.name, - False, - ) - - raw_map = _type_helpers.unpack_literal_map_to_sdk_python_std( - _literal_models.LiteralMap.from_flyte_idl( - _utils.load_proto_from_file( - _literals_pb2.LiteralMap, - os.path.join(input_dir, _constants.OUTPUT_FILE_NAME), - ) - ), - _type_map_from_variable_map(_task_defs.add_one.interface.outputs), - ) - assert raw_map["b"] == 10 - assert len(raw_map) == 1 - - # reset the env vars - if orig_env_index_var_name: - os.environ["BATCH_JOB_ARRAY_INDEX_VAR_NAME"] = orig_env_index_var_name - if orig_env_array_index: - os.environ["AWS_BATCH_JOB_ARRAY_INDEX"] = orig_env_array_index - - -@mock.patch("flytekit.bin.entrypoint._legacy_execute_task") -def test_backwards_compatible_replacement(mock_legacy_execute_task): - def return_args(*args, **kwargs): - assert args[4] is None - - mock_legacy_execute_task.side_effect = return_args - - with _TemporaryConfiguration( - os.path.join(os.path.dirname(__file__), "fake.config"), - internal_overrides={"project": "test", "domain": "development"}, - ): - with _utils.AutoDeletingTempDir("in"): - with _utils.AutoDeletingTempDir("out"): - cmd = [] - cmd.extend(["--task-module", "fake"]) - cmd.extend(["--task-name", "fake"]) - cmd.extend(["--inputs", "fake"]) - cmd.extend(["--output-prefix", "fake"]) - cmd.extend(["--raw-output-data-prefix", "{{.rawOutputDataPrefix}}"]) - result = CliRunner().invoke(execute_task_cmd, cmd) - assert result.exit_code == 0 - - -@mock.patch("flytekit.common.utils.load_proto_from_file") + + +@mock.patch("flytekit.core.utils.load_proto_from_file") @mock.patch("flytekit.core.data_persistence.FileAccessProvider.get_data") @mock.patch("flytekit.core.data_persistence.FileAccessProvider.put_data") -@mock.patch("flytekit.common.utils.write_proto_to_file") +@mock.patch("flytekit.core.utils.write_proto_to_file") def test_dispatch_execute_void(mock_write_to_file, mock_upload_dir, mock_get_data, mock_load_proto): # Just leave these here, mock them out so nothing happens mock_get_data.return_value = True @@ -214,10 +50,10 @@ def verify_output(*args, **kwargs): assert mock_write_to_file.call_count == 1 -@mock.patch("flytekit.common.utils.load_proto_from_file") +@mock.patch("flytekit.core.utils.load_proto_from_file") @mock.patch("flytekit.core.data_persistence.FileAccessProvider.get_data") @mock.patch("flytekit.core.data_persistence.FileAccessProvider.put_data") -@mock.patch("flytekit.common.utils.write_proto_to_file") +@mock.patch("flytekit.core.utils.write_proto_to_file") def test_dispatch_execute_ignore(mock_write_to_file, mock_upload_dir, mock_get_data, mock_load_proto): # Just leave these here, mock them out so nothing happens mock_get_data.return_value = True @@ -243,10 +79,10 @@ def test_dispatch_execute_ignore(mock_write_to_file, mock_upload_dir, mock_get_d assert mock_write_to_file.call_count == 0 -@mock.patch("flytekit.common.utils.load_proto_from_file") +@mock.patch("flytekit.core.utils.load_proto_from_file") @mock.patch("flytekit.core.data_persistence.FileAccessProvider.get_data") @mock.patch("flytekit.core.data_persistence.FileAccessProvider.put_data") -@mock.patch("flytekit.common.utils.write_proto_to_file") +@mock.patch("flytekit.core.utils.write_proto_to_file") def test_dispatch_execute_exception(mock_write_to_file, mock_upload_dir, mock_get_data, mock_load_proto): # Just leave these here, mock them out so nothing happens mock_get_data.return_value = True @@ -273,7 +109,7 @@ def verify_output(*args, **kwargs): # This function collects outputs instead of writing them to a file. -# See flytekit.common.utils.write_proto_to_file for the original +# See flytekit.core.utils.write_proto_to_file for the original def get_output_collector(results: OrderedDict): def output_collector(proto, path): results[path] = proto @@ -281,10 +117,10 @@ def output_collector(proto, path): return output_collector -@mock.patch("flytekit.common.utils.load_proto_from_file") +@mock.patch("flytekit.core.utils.load_proto_from_file") @mock.patch("flytekit.core.data_persistence.FileAccessProvider.get_data") @mock.patch("flytekit.core.data_persistence.FileAccessProvider.put_data") -@mock.patch("flytekit.common.utils.write_proto_to_file") +@mock.patch("flytekit.core.utils.write_proto_to_file") def test_dispatch_execute_normal(mock_write_to_file, mock_upload_dir, mock_get_data, mock_load_proto): # Just leave these here, mock them out so nothing happens mock_get_data.return_value = True @@ -318,10 +154,10 @@ def t1(a: int) -> str: assert lm.literals["o0"].scalar.primitive.string_value == "string is: 5" -@mock.patch("flytekit.common.utils.load_proto_from_file") +@mock.patch("flytekit.core.utils.load_proto_from_file") @mock.patch("flytekit.core.data_persistence.FileAccessProvider.get_data") @mock.patch("flytekit.core.data_persistence.FileAccessProvider.put_data") -@mock.patch("flytekit.common.utils.write_proto_to_file") +@mock.patch("flytekit.core.utils.write_proto_to_file") def test_dispatch_execute_user_error_non_recov(mock_write_to_file, mock_upload_dir, mock_get_data, mock_load_proto): # Just leave these here, mock them out so nothing happens mock_get_data.return_value = True @@ -358,10 +194,10 @@ def t1(a: int) -> str: assert ed.error.origin == execution_models.ExecutionError.ErrorKind.USER -@mock.patch("flytekit.common.utils.load_proto_from_file") +@mock.patch("flytekit.core.utils.load_proto_from_file") @mock.patch("flytekit.core.data_persistence.FileAccessProvider.get_data") @mock.patch("flytekit.core.data_persistence.FileAccessProvider.put_data") -@mock.patch("flytekit.common.utils.write_proto_to_file") +@mock.patch("flytekit.core.utils.write_proto_to_file") def test_dispatch_execute_user_error_recoverable(mock_write_to_file, mock_upload_dir, mock_get_data, mock_load_proto): # Just leave these here, mock them out so nothing happens mock_get_data.return_value = True @@ -402,10 +238,10 @@ def my_subwf(a: int) -> typing.List[str]: assert ed.error.origin == execution_models.ExecutionError.ErrorKind.USER -@mock.patch("flytekit.common.utils.load_proto_from_file") +@mock.patch("flytekit.core.utils.load_proto_from_file") @mock.patch("flytekit.core.data_persistence.FileAccessProvider.get_data") @mock.patch("flytekit.core.data_persistence.FileAccessProvider.put_data") -@mock.patch("flytekit.common.utils.write_proto_to_file") +@mock.patch("flytekit.core.utils.write_proto_to_file") def test_dispatch_execute_system_error(mock_write_to_file, mock_upload_dir, mock_get_data, mock_load_proto): # Just leave these here, mock them out so nothing happens mock_get_data.return_value = True diff --git a/tests/flytekit/unit/cli/pyflyte/test_launch_plans.py b/tests/flytekit/unit/cli/pyflyte/test_launch_plans.py deleted file mode 100644 index e19fea82a0..0000000000 --- a/tests/flytekit/unit/cli/pyflyte/test_launch_plans.py +++ /dev/null @@ -1,30 +0,0 @@ -import pytest - -from flytekit.clis.sdk_in_container import launch_plan -from flytekit.clis.sdk_in_container.launch_plan import launch_plans - - -def test_list_commands(mock_ctx): - g = launch_plan.LaunchPlanExecuteGroup("test_group") - v = g.list_commands(mock_ctx) - assert v == ["common.workflows.simple.SimpleWorkflow"] - - -def test_get_commands(mock_ctx): - g = launch_plan.LaunchPlanExecuteGroup("test_group") - v = g.get_command(mock_ctx, "common.workflows.simple.SimpleWorkflow") - assert v.params[0].human_readable_name == "input_1" - assert "INTEGER" in v.params[0].help - assert v.params[1].human_readable_name == "input_2" - assert "INTEGER" in v.params[1].help - assert "Not required." in v.params[1].help - - with pytest.raises(Exception): - g.get_command(mock_ctx, "common.workflows.simple.DoesNotExist") - with pytest.raises(Exception): - g.get_command(mock_ctx, "does.not.exist") - - -def test_launch_plans_commands(mock_ctx): - command_names = [c for c in launch_plans.list_commands(mock_ctx)] - assert command_names == sorted(["execute", "activate-all", "activate-all-schedules"]) diff --git a/tests/flytekit/unit/cli/pyflyte/test_package.py b/tests/flytekit/unit/cli/pyflyte/test_package.py index e14be7fc5b..ced8c716c5 100644 --- a/tests/flytekit/unit/cli/pyflyte/test_package.py +++ b/tests/flytekit/unit/cli/pyflyte/test_package.py @@ -7,8 +7,8 @@ import flytekit from flytekit.clis.sdk_in_container import package, pyflyte, serialize -from flytekit.common.exceptions.user import FlyteValidationException from flytekit.core import context_manager +from flytekit.exceptions.user import FlyteValidationException def test_validate_image(): diff --git a/tests/flytekit/unit/cli/pyflyte/test_register.py b/tests/flytekit/unit/cli/pyflyte/test_register.py deleted file mode 100644 index d1a6eed495..0000000000 --- a/tests/flytekit/unit/cli/pyflyte/test_register.py +++ /dev/null @@ -1,40 +0,0 @@ -from mock import MagicMock - -from flytekit.common.launch_plan import SdkLaunchPlan -from flytekit.common.tasks.task import SdkTask -from flytekit.common.workflow import SdkWorkflow - - -def test_register_workflows(mock_clirunner, monkeypatch): - - mock_register_task = MagicMock(return_value=MagicMock()) - monkeypatch.setattr(SdkTask, "register", mock_register_task) - mock_register_workflow = MagicMock(return_value=MagicMock()) - monkeypatch.setattr(SdkWorkflow, "register", mock_register_workflow) - mock_register_launch_plan = MagicMock(return_value=MagicMock()) - monkeypatch.setattr(SdkLaunchPlan, "register", mock_register_launch_plan) - - result = mock_clirunner("register", "-p", "project", "-d", "development", "-v", "--version", "workflows") - - assert result.exit_code == 0 - - assert len(mock_register_task.mock_calls) == 4 - assert len(mock_register_workflow.mock_calls) == 1 - assert len(mock_register_launch_plan.mock_calls) == 1 - - -def test_register_workflows_with_test_switch(mock_clirunner, monkeypatch): - mock_register_task = MagicMock(return_value=MagicMock()) - monkeypatch.setattr(SdkTask, "register", mock_register_task) - mock_register_workflow = MagicMock(return_value=MagicMock()) - monkeypatch.setattr(SdkWorkflow, "register", mock_register_workflow) - mock_register_launch_plan = MagicMock(return_value=MagicMock()) - monkeypatch.setattr(SdkLaunchPlan, "register", mock_register_launch_plan) - - result = mock_clirunner("register", "-p", "project", "-d", "development", "-v", "--version", "--test", "workflows") - - assert result.exit_code == 0 - - assert len(mock_register_task.mock_calls) == 0 - assert len(mock_register_workflow.mock_calls) == 0 - assert len(mock_register_launch_plan.mock_calls) == 0 diff --git a/tests/flytekit/unit/cli/test_cli_helpers.py b/tests/flytekit/unit/cli/test_cli_helpers.py index 9188b2fc3a..3ed08848c3 100644 --- a/tests/flytekit/unit/cli/test_cli_helpers.py +++ b/tests/flytekit/unit/cli/test_cli_helpers.py @@ -2,15 +2,12 @@ import flyteidl.admin.task_pb2 as _task_pb2 import flyteidl.admin.workflow_pb2 as _workflow_pb2 import flyteidl.core.tasks_pb2 as _core_task_pb2 -import pytest from flyteidl.core import identifier_pb2 as _identifier_pb2 from flyteidl.core import workflow_pb2 as _core_workflow_pb2 from flyteidl.core.identifier_pb2 import LAUNCH_PLAN from flytekit.clis import helpers from flytekit.clis.helpers import _hydrate_identifier, _hydrate_workflow_template_nodes, hydrate_registration_parameters -from flytekit.models import literals, types -from flytekit.models.interface import Parameter, ParameterMap, Variable def test_parse_args_into_dict(): @@ -28,36 +25,6 @@ def test_parse_args_into_dict(): assert output == {} -def test_construct_literal_map_from_variable_map(): - v = Variable(type=types.LiteralType(simple=types.SimpleType.INTEGER), description="some description") - variable_map = { - "inputa": v, - } - - input_txt_dictionary = {"inputa": "15"} - - literal_map = helpers.construct_literal_map_from_variable_map(variable_map, input_txt_dictionary) - parsed_literal = literal_map.literals["inputa"].value - ll = literals.Scalar(primitive=literals.Primitive(integer=15)) - assert parsed_literal == ll - - -def test_construct_literal_map_from_parameter_map(): - v = Variable(type=types.LiteralType(simple=types.SimpleType.INTEGER), description="some description") - p = Parameter(var=v, required=True) - pm = ParameterMap(parameters={"inputa": p}) - - input_txt_dictionary = {"inputa": "15"} - - literal_map = helpers.construct_literal_map_from_parameter_map(pm, input_txt_dictionary) - parsed_literal = literal_map.literals["inputa"].value - ll = literals.Scalar(primitive=literals.Primitive(integer=15)) - assert parsed_literal == ll - - with pytest.raises(Exception): - helpers.construct_literal_map_from_parameter_map(pm, {}) - - def test_strtobool(): assert not helpers.str2bool("False") assert not helpers.str2bool("OFF") diff --git a/tests/flytekit/unit/cli/test_flyte_cli.py b/tests/flytekit/unit/cli/test_flyte_cli.py index 3049590036..7c5830b728 100644 --- a/tests/flytekit/unit/cli/test_flyte_cli.py +++ b/tests/flytekit/unit/cli/test_flyte_cli.py @@ -4,54 +4,17 @@ from click.testing import CliRunner as _CliRunner from flytekit.clis.flyte_cli import main as _main -from flytekit.common.exceptions.user import FlyteAssertion -from flytekit.common.types import primitives -from flytekit.configuration import TemporaryConfiguration +from flytekit.exceptions.user import FlyteAssertion from flytekit.models import filters as _filters from flytekit.models.admin import common as _admin_common from flytekit.models.core import identifier as _core_identifier from flytekit.models.project import Project as _Project -from flytekit.sdk.tasks import inputs, outputs, python_task mm = _mock.MagicMock() mm.return_value = 100 -def get_sample_task(): - """ - :rtype: flytekit.common.tasks.task.SdkTask - """ - - @inputs(a=primitives.Integer) - @outputs(b=primitives.Integer) - @python_task() - def my_task(wf_params, a, b): - b.set(a + 1) - - return my_task - - -@_mock.patch("flytekit.clis.flyte_cli.main._load_proto_from_file") -def test__extract_files(load_mock): - t = get_sample_task() - with TemporaryConfiguration( - "", - internal_overrides={"image": "myflyteimage:v123", "project": "myflyteproject", "domain": "development"}, - ): - task_spec = t.serialize() - - load_mock.side_effect = [task_spec] - new_id, entity = _main._extract_pair("a", 1, "myproject", "development", "v", {}) - assert ( - new_id - == _core_identifier.Identifier( - _core_identifier.ResourceType.TASK, "myproject", "development", "test_flyte_cli.my_task", "v" - ).to_flyte_idl() - ) - assert task_spec == entity - - -@_mock.patch("flytekit.clis.flyte_cli.main._load_proto_from_file") +@_mock.patch("flytekit.clis.flyte_cli.main.utils") def test__extract_files_with_unspecified_resource_type(load_mock): id = _core_identifier.Identifier( _core_identifier.ResourceType.UNSPECIFIED, @@ -61,7 +24,7 @@ def test__extract_files_with_unspecified_resource_type(load_mock): "v", ) - load_mock.return_value = id.to_flyte_idl() + load_mock.load_proto_from_file.return_value = id.to_flyte_idl() with pytest.raises(FlyteAssertion): _main._extract_pair("a", "b", "myflyteproject", "development", "v", {}) diff --git a/tests/flytekit/unit/common_tests/__init__.py b/tests/flytekit/unit/common_tests/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/tests/flytekit/unit/common_tests/exceptions/__init__.py b/tests/flytekit/unit/common_tests/exceptions/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/tests/flytekit/unit/common_tests/mixins/__init__.py b/tests/flytekit/unit/common_tests/mixins/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/tests/flytekit/unit/common_tests/mixins/sample_registerable.py b/tests/flytekit/unit/common_tests/mixins/sample_registerable.py deleted file mode 100644 index b0ecb08ca9..0000000000 --- a/tests/flytekit/unit/common_tests/mixins/sample_registerable.py +++ /dev/null @@ -1,15 +0,0 @@ -from flytekit.common import sdk_bases as _sdk_bases -from flytekit.common.mixins import registerable as _registerable - - -class ExampleRegisterable( - _registerable.RegisterableEntity, _registerable.TrackableEntity, metaclass=_sdk_bases.ExtendedSdkType -): - def __init__(self, *args, **kwargs): - super(ExampleRegisterable, self).__init__(*args, **kwargs) - - def promote_from_model(cls, base_model): - pass - - -example = ExampleRegisterable() diff --git a/tests/flytekit/unit/common_tests/mixins/test_registerable.py b/tests/flytekit/unit/common_tests/mixins/test_registerable.py deleted file mode 100644 index 030d5bc1e9..0000000000 --- a/tests/flytekit/unit/common_tests/mixins/test_registerable.py +++ /dev/null @@ -1,13 +0,0 @@ -from tests.flytekit.unit.common_tests.mixins import sample_registerable as _sample_registerable - - -def test_instance_tracker(): - assert _sample_registerable.example.instantiated_in == "tests.flytekit.unit.common_tests.mixins.sample_registerable" - - -def test_auto_name_assignment(): - _sample_registerable.example.auto_assign_name() - assert ( - _sample_registerable.example.platform_valid_name - == "tests.flytekit.unit.common_tests.mixins.sample_registerable.example" - ) diff --git a/tests/flytekit/unit/common_tests/tasks/__init__.py b/tests/flytekit/unit/common_tests/tasks/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/tests/flytekit/unit/common_tests/tasks/spark/__init__.py b/tests/flytekit/unit/common_tests/tasks/spark/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/tests/flytekit/unit/common_tests/tasks/spark/test_spark_task.py b/tests/flytekit/unit/common_tests/tasks/spark/test_spark_task.py deleted file mode 100644 index 9389fdc7ea..0000000000 --- a/tests/flytekit/unit/common_tests/tasks/spark/test_spark_task.py +++ /dev/null @@ -1,37 +0,0 @@ -from six.moves import range - -from flytekit.sdk.tasks import outputs, spark_task -from flytekit.sdk.types import Types - -# This file is in a subdirectory to make it easier to exclude when not running in a container -# and pyspark is not available - - -@outputs(out=Types.Integer) -@spark_task(retries=1) -def my_spark_task(wf, sc, out): - def _inside(p): - return p < 1000 - - count = sc.parallelize(range(0, 10000)).filter(_inside).count() - out.set(count) - - -@outputs(out=Types.Integer) -@spark_task(retries=3) -def my_spark_task2(wf, sc, out): - # This test makes sure spark_task doesn't choke on a non-package module and modules which overlap with auto-included - # modules. - def _inside(p): - return p < 500 - - count = sc.parallelize(range(0, 10000)).filter(_inside).count() - out.set(count) - - -def test_basic_spark_execution(): - outputs = my_spark_task.unit_test() - assert outputs["out"] == 1000 - - outputs = my_spark_task2.unit_test() - assert outputs["out"] == 500 diff --git a/tests/flytekit/unit/common_tests/tasks/test_execution_params.py b/tests/flytekit/unit/common_tests/tasks/test_execution_params.py deleted file mode 100644 index cf634bfa96..0000000000 --- a/tests/flytekit/unit/common_tests/tasks/test_execution_params.py +++ /dev/null @@ -1,76 +0,0 @@ -import os - -import py -import pytest - -from flytekit.common.tasks.sdk_runnable import SecretsManager -from flytekit.configuration import secrets - - -def test_secrets_manager_default(): - with pytest.raises(ValueError): - sec = SecretsManager() - sec.get("group", "key") - - -def test_secrets_manager_get_envvar(): - sec = SecretsManager() - with pytest.raises(ValueError): - sec.get_secrets_env_var("test", "") - with pytest.raises(ValueError): - sec.get_secrets_env_var("", "x") - assert sec.get_secrets_env_var("group", "test") == f"{secrets.SECRETS_ENV_PREFIX.get()}GROUP_TEST" - - -def test_secrets_manager_get_file(): - sec = SecretsManager() - with pytest.raises(ValueError): - sec.get_secrets_file("test", "") - with pytest.raises(ValueError): - sec.get_secrets_file("", "x") - assert sec.get_secrets_file("group", "test") == os.path.join( - secrets.SECRETS_DEFAULT_DIR.get(), - "group", - f"{secrets.SECRETS_FILE_PREFIX.get()}test", - ) - - -def test_secrets_manager_file(tmpdir: py.path.local): - tmp = tmpdir.mkdir("file_test").dirname - os.environ["FLYTE_SECRETS_DEFAULT_DIR"] = tmp - sec = SecretsManager() - f = os.path.join(tmp, "test") - with open(f, "w+") as w: - w.write("my-password") - - with pytest.raises(ValueError): - sec.get("test", "") - with pytest.raises(ValueError): - sec.get("", "x") - # Group dir not exists - with pytest.raises(ValueError): - sec.get("group", "test") - - g = os.path.join(tmp, "group") - os.makedirs(g) - f = os.path.join(g, "test") - with open(f, "w+") as w: - w.write("my-password") - assert sec.get("group", "test") == "my-password" - del os.environ["FLYTE_SECRETS_DEFAULT_DIR"] - - -def test_secrets_manager_bad_env(): - with pytest.raises(ValueError): - os.environ["TEST"] = "value" - sec = SecretsManager() - sec.get("group", "test") - - -def test_secrets_manager_env(): - sec = SecretsManager() - os.environ[sec.get_secrets_env_var("group", "test")] = "value" - assert sec.get("group", "test") == "value" - - os.environ[sec.get_secrets_env_var(group="group", key="key")] = "value" - assert sec.get(group="group", key="key") == "value" diff --git a/tests/flytekit/unit/common_tests/tasks/test_raw_container_task.py b/tests/flytekit/unit/common_tests/tasks/test_raw_container_task.py deleted file mode 100644 index 1267ce1c8b..0000000000 --- a/tests/flytekit/unit/common_tests/tasks/test_raw_container_task.py +++ /dev/null @@ -1,27 +0,0 @@ -from flytekit.common.tasks.raw_container import SdkRawContainerTask -from flytekit.sdk.types import Types - - -def test_raw_container_task_definition(): - tk = SdkRawContainerTask( - inputs={"x": Types.Integer}, - outputs={"y": Types.Integer}, - image="my-image", - command=["echo", "hello, world!"], - gpu_limit="1", - gpu_request="1", - ) - assert not tk.serialize() is None - - -def test_raw_container_task_definition_no_outputs(): - tk = SdkRawContainerTask( - inputs={"x": Types.Integer}, - image="my-image", - command=["echo", "hello, world!"], - gpu_limit="1", - gpu_request="1", - ) - assert not tk.serialize() is None - task_instance = tk(x=3) - assert task_instance.inputs[0].binding.scalar.primitive.integer == 3 diff --git a/tests/flytekit/unit/common_tests/tasks/test_sdk_runnable.py b/tests/flytekit/unit/common_tests/tasks/test_sdk_runnable.py deleted file mode 100644 index 43e677551c..0000000000 --- a/tests/flytekit/unit/common_tests/tasks/test_sdk_runnable.py +++ /dev/null @@ -1,44 +0,0 @@ -import pytest as _pytest - -from flytekit.common import constants as _common_constants -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.tasks import sdk_runnable -from flytekit.common.types import primitives -from flytekit.models import interface - - -def test_basic_unit_test(): - def add_one(wf_params, value_in, value_out): - value_out.set(value_in + 1) - - t = sdk_runnable.SdkRunnableTask( - add_one, - _common_constants.SdkTaskType.PYTHON_TASK, - "1", - 1, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - False, - None, - {}, - False, - None, - ) - t.add_inputs({"value_in": interface.Variable(primitives.Integer.to_flyte_literal_type(), "")}) - t.add_outputs({"value_out": interface.Variable(primitives.Integer.to_flyte_literal_type(), "")}) - out = t.unit_test(value_in=1) - assert out["value_out"] == 2 - - with _pytest.raises(_user_exceptions.FlyteAssertion) as e: - t() - - assert "value_in" in str(e.value) - assert "INTEGER" in str(e.value) diff --git a/tests/flytekit/unit/common_tests/tasks/test_task.py b/tests/flytekit/unit/common_tests/tasks/test_task.py deleted file mode 100644 index e33a757412..0000000000 --- a/tests/flytekit/unit/common_tests/tasks/test_task.py +++ /dev/null @@ -1,108 +0,0 @@ -import os as _os - -import pytest as _pytest -from flyteidl.admin import task_pb2 as _admin_task_pb2 -from mock import MagicMock as _MagicMock -from mock import patch as _patch - -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.tasks import task as _task -from flytekit.common.tasks.presto_task import SdkPrestoTask -from flytekit.common.types import primitives -from flytekit.configuration import TemporaryConfiguration -from flytekit.models import task as _task_models -from flytekit.models.core import identifier as _identifier -from flytekit.sdk.tasks import inputs, outputs, python_task -from flytekit.sdk.types import Types - - -@_patch("flytekit.engines.flyte.engine._FlyteClientManager") -@_patch("flytekit.configuration.platform.URL") -def test_fetch_latest(mock_url, mock_client_manager): - mock_url.get.return_value = "localhost" - admin_task = _task_models.Task( - _identifier.Identifier(_identifier.ResourceType.TASK, "p1", "d1", "n1", "v1"), - _MagicMock(), - ) - mock_client = _MagicMock() - mock_client.list_tasks_paginated = _MagicMock(return_value=([admin_task], "")) - mock_client_manager.return_value.client = mock_client - task = _task.SdkTask.fetch_latest("p1", "d1", "n1") - assert task.id == admin_task.id - - -@_patch("flytekit.engines.flyte.engine._FlyteClientManager") -@_patch("flytekit.configuration.platform.URL") -def test_fetch_latest_not_exist(mock_url, mock_client_manager): - mock_client = _MagicMock() - mock_client.list_tasks_paginated = _MagicMock(return_value=(None, "")) - mock_client_manager.return_value.client = mock_client - mock_url.get.return_value = "localhost" - with _pytest.raises(_user_exceptions.FlyteEntityNotExistException): - _task.SdkTask.fetch_latest("p1", "d1", "n1") - - -def get_sample_task(): - """ - :rtype: flytekit.common.tasks.task.SdkTask - """ - - @inputs(a=primitives.Integer) - @outputs(b=primitives.Integer) - @python_task() - def my_task(wf_params, a, b): - b.set(a + 1) - - return my_task - - -def test_task_serialization(): - t = get_sample_task() - with TemporaryConfiguration( - _os.path.join( - _os.path.dirname(_os.path.realpath(__file__)), - "../../../common/configs/local.config", - ), - internal_overrides={"image": "myflyteimage:v123", "project": "myflyteproject", "domain": "development"}, - ): - s = t.serialize() - - assert isinstance(s, _admin_task_pb2.TaskSpec) - assert s.template.id.name == "tests.flytekit.unit.common_tests.tasks.test_task.my_task" - assert s.template.container.image == "myflyteimage:v123" - - -schema = Types.Schema([("a", Types.String), ("b", Types.Integer)]) - - -def test_task_produce_deterministic_version(): - containerless_task = SdkPrestoTask( - task_inputs=inputs(ds=Types.String, rg=Types.String), - statement="SELECT * FROM flyte.widgets WHERE ds = '{{ .Inputs.ds}}' LIMIT 10", - output_schema=schema, - routing_group="{{ .Inputs.rg }}", - ) - identical_containerless_task = SdkPrestoTask( - task_inputs=inputs(ds=Types.String, rg=Types.String), - statement="SELECT * FROM flyte.widgets WHERE ds = '{{ .Inputs.ds}}' LIMIT 10", - output_schema=schema, - routing_group="{{ .Inputs.rg }}", - ) - different_containerless_task = SdkPrestoTask( - task_inputs=inputs(ds=Types.String, rg=Types.String), - statement="SELECT * FROM flyte.widgets WHERE ds = '{{ .Inputs.ds}}' LIMIT 100000", - output_schema=schema, - routing_group="{{ .Inputs.rg }}", - ) - assert ( - containerless_task._produce_deterministic_version() - == identical_containerless_task._produce_deterministic_version() - ) - - assert ( - containerless_task._produce_deterministic_version() - != different_containerless_task._produce_deterministic_version() - ) - - with _pytest.raises(Exception): - get_sample_task()._produce_deterministic_version() diff --git a/tests/flytekit/unit/common_tests/test_interface.py b/tests/flytekit/unit/common_tests/test_interface.py deleted file mode 100644 index b6627c1a6b..0000000000 --- a/tests/flytekit/unit/common_tests/test_interface.py +++ /dev/null @@ -1,80 +0,0 @@ -import pytest - -from flytekit.common import interface -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.types import containers, primitives - - -def test_binding_data_primitive_static(): - upstream_nodes = set() - bd = interface.BindingData.from_python_std( - primitives.Float.to_flyte_literal_type(), 3.0, upstream_nodes=upstream_nodes - ) - - assert len(upstream_nodes) == 0 - assert bd.promise is None - assert bd.collection is None - assert bd.map is None - assert bd.scalar.primitive.float_value == 3.0 - - assert interface.BindingData.from_flyte_idl(bd.to_flyte_idl()) == bd - - with pytest.raises(_user_exceptions.FlyteTypeException): - interface.BindingData.from_python_std( - primitives.Float.to_flyte_literal_type(), - "abc", - ) - - with pytest.raises(_user_exceptions.FlyteTypeException): - interface.BindingData.from_python_std( - primitives.Float.to_flyte_literal_type(), - [1.0, 2.0, 3.0], - ) - - -def test_binding_data_list_static(): - upstream_nodes = set() - bd = interface.BindingData.from_python_std( - containers.List(primitives.String).to_flyte_literal_type(), - ["abc", "cde"], - upstream_nodes=upstream_nodes, - ) - - assert len(upstream_nodes) == 0 - assert bd.promise is None - assert bd.collection.bindings[0].scalar.primitive.string_value == "abc" - assert bd.collection.bindings[1].scalar.primitive.string_value == "cde" - assert bd.map is None - assert bd.scalar is None - - assert interface.BindingData.from_flyte_idl(bd.to_flyte_idl()) == bd - - with pytest.raises(_user_exceptions.FlyteTypeException): - interface.BindingData.from_python_std( - containers.List(primitives.String).to_flyte_literal_type(), - "abc", - ) - - with pytest.raises(_user_exceptions.FlyteTypeException): - interface.BindingData.from_python_std( - containers.List(primitives.String).to_flyte_literal_type(), [1.0, 2.0, 3.0] - ) - - -def test_binding_generic_map_static(): - upstream_nodes = set() - bd = interface.BindingData.from_python_std( - primitives.Generic.to_flyte_literal_type(), - {"a": "hi", "b": [1, 2, 3], "c": {"d": "e"}}, - upstream_nodes=upstream_nodes, - ) - - assert len(upstream_nodes) == 0 - assert bd.promise is None - assert bd.map is None - assert bd.scalar.generic["a"] == "hi" - assert bd.scalar.generic["b"].values[0].number_value == 1.0 - assert bd.scalar.generic["b"].values[1].number_value == 2.0 - assert bd.scalar.generic["b"].values[2].number_value == 3.0 - assert bd.scalar.generic["c"]["d"] == "e" - assert interface.BindingData.from_flyte_idl(bd.to_flyte_idl()) == bd diff --git a/tests/flytekit/unit/common_tests/test_launch_plan.py b/tests/flytekit/unit/common_tests/test_launch_plan.py deleted file mode 100644 index 92943d37f1..0000000000 --- a/tests/flytekit/unit/common_tests/test_launch_plan.py +++ /dev/null @@ -1,388 +0,0 @@ -import os as _os - -import pytest as _pytest - -from flytekit import configuration as _configuration -from flytekit.common import launch_plan as _launch_plan -from flytekit.common import notifications as _notifications -from flytekit.common import schedules as _schedules -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.models import common as _common_models -from flytekit.models import schedule as _schedule -from flytekit.models import types as _type_models -from flytekit.models.core import execution as _execution -from flytekit.models.core import identifier as _identifier -from flytekit.sdk import types as _types -from flytekit.sdk import workflow as _workflow - - -def test_default_assumable_iam_role(): - with _configuration.TemporaryConfiguration( - _os.path.join( - _os.path.dirname(_os.path.realpath(__file__)), - "../../common/configs/local.config", - ) - ): - workflow_to_test = _workflow.workflow( - {}, - inputs={ - "required_input": _workflow.Input(_types.Types.Integer), - "default_input": _workflow.Input(_types.Types.Integer, default=5), - }, - ) - lp = workflow_to_test.create_launch_plan() - assert lp.auth_role.assumable_iam_role == "arn:aws:iam::ABC123:role/my-flyte-role" - - -def test_hard_coded_assumable_iam_role(): - workflow_to_test = _workflow.workflow( - {}, - inputs={ - "required_input": _workflow.Input(_types.Types.Integer), - "default_input": _workflow.Input(_types.Types.Integer, default=5), - }, - ) - lp = workflow_to_test.create_launch_plan(assumable_iam_role="override") - assert lp.auth_role.assumable_iam_role == "override" - - -def test_default_deprecated_role(): - with _configuration.TemporaryConfiguration( - _os.path.join( - _os.path.dirname(_os.path.realpath(__file__)), - "../../common/configs/deprecated_local.config", - ) - ): - workflow_to_test = _workflow.workflow( - {}, - inputs={ - "required_input": _workflow.Input(_types.Types.Integer), - "default_input": _workflow.Input(_types.Types.Integer, default=5), - }, - ) - lp = workflow_to_test.create_launch_plan() - assert lp.auth_role.assumable_iam_role == "arn:aws:iam::ABC123:role/my-flyte-role" - - -def test_hard_coded_deprecated_role(): - workflow_to_test = _workflow.workflow( - {}, - inputs={ - "required_input": _workflow.Input(_types.Types.Integer), - "default_input": _workflow.Input(_types.Types.Integer, default=5), - }, - ) - lp = workflow_to_test.create_launch_plan(role="override") - assert lp.auth_role.assumable_iam_role == "override" - - -def test_kubernetes_service_account(): - workflow_to_test = _workflow.workflow( - {}, - inputs={ - "required_input": _workflow.Input(_types.Types.Integer), - "default_input": _workflow.Input(_types.Types.Integer, default=5), - }, - ) - lp = workflow_to_test.create_launch_plan(kubernetes_service_account="kube-service-acct") - assert lp.auth_role.kubernetes_service_account == "kube-service-acct" - - -def test_fixed_inputs(): - workflow_to_test = _workflow.workflow( - {}, - inputs={ - "required_input": _workflow.Input(_types.Types.Integer), - "default_input": _workflow.Input(_types.Types.Integer, default=5), - }, - ) - lp = workflow_to_test.create_launch_plan(fixed_inputs={"required_input": 4}) - assert len(lp.fixed_inputs.literals) == 1 - assert lp.fixed_inputs.literals["required_input"].scalar.primitive.integer == 4 - assert len(lp.default_inputs.parameters) == 1 - assert lp.default_inputs.parameters["default_input"].default.scalar.primitive.integer == 5 - - -def test_redefining_inputs_good(): - workflow_to_test = _workflow.workflow( - {}, - inputs={ - "required_input": _workflow.Input(_types.Types.Integer), - "default_input": _workflow.Input(_types.Types.Integer, default=5), - }, - ) - lp = workflow_to_test.create_launch_plan( - default_inputs={"required_input": _workflow.Input(_types.Types.Integer, default=900)} - ) - assert len(lp.fixed_inputs.literals) == 0 - assert len(lp.default_inputs.parameters) == 2 - assert lp.default_inputs.parameters["required_input"].default.scalar.primitive.integer == 900 - assert lp.default_inputs.parameters["default_input"].default.scalar.primitive.integer == 5 - - -def test_no_additional_inputs(): - workflow_to_test = _workflow.workflow( - {}, - inputs={ - "required_input": _workflow.Input(_types.Types.Integer), - "default_input": _workflow.Input(_types.Types.Integer, default=5), - }, - ) - lp = workflow_to_test.create_launch_plan() - assert len(lp.fixed_inputs.literals) == 0 - assert lp.default_inputs.parameters["default_input"].default.scalar.primitive.integer == 5 - assert lp.default_inputs.parameters["required_input"].required is True - - -@_pytest.mark.parametrize( - "schedule,cron_expression,cron_schedule", - [ - (_schedules.CronSchedule("* * ? * * *"), "* * ? * * *", None), - (_schedules.CronSchedule(cron_expression="* * ? * * *"), "* * ? * * *", None), - (_schedules.CronSchedule(cron_expression="0/15 * * * ? *"), "0/15 * * * ? *", None), - (_schedules.CronSchedule(schedule="* * * * *"), None, _schedule.Schedule.CronSchedule("* * * * *", None)), - ( - _schedules.CronSchedule(schedule="* * * * *", offset="P1D"), - None, - _schedule.Schedule.CronSchedule("* * * * *", "P1D"), - ), - ], -) -def test_schedule(schedule, cron_expression, cron_schedule): - workflow_to_test = _workflow.workflow( - {}, - inputs={ - "required_input": _workflow.Input(_types.Types.Integer), - "default_input": _workflow.Input(_types.Types.Integer, default=5), - }, - ) - lp = workflow_to_test.create_launch_plan( - fixed_inputs={"required_input": 5}, - schedule=schedule, - role="what", - ) - assert lp.entity_metadata.schedule.kickoff_time_input_arg is None - assert lp.entity_metadata.schedule.cron_expression == cron_expression - assert lp.entity_metadata.schedule.cron_schedule == cron_schedule - assert lp.is_scheduled - - -def test_no_schedule(): - workflow_to_test = _workflow.workflow( - {}, - inputs={ - "required_input": _workflow.Input(_types.Types.Integer), - "default_input": _workflow.Input(_types.Types.Integer, default=5), - }, - ) - lp = workflow_to_test.create_launch_plan() - assert lp.entity_metadata.schedule.kickoff_time_input_arg == "" - assert lp.entity_metadata.schedule.schedule_expression is None - assert not lp.is_scheduled - - -def test_schedule_pointing_to_datetime(): - workflow_to_test = _workflow.workflow( - {}, - inputs={ - "required_input": _workflow.Input(_types.Types.Datetime), - "default_input": _workflow.Input(_types.Types.Integer, default=5), - }, - ) - lp = workflow_to_test.create_launch_plan( - schedule=_schedules.CronSchedule("* * ? * * *", kickoff_time_input_arg="required_input"), - role="what", - ) - assert lp.entity_metadata.schedule.kickoff_time_input_arg == "required_input" - assert lp.entity_metadata.schedule.cron_expression == "* * ? * * *" - - -def test_notifications(): - workflow_to_test = _workflow.workflow( - {}, - inputs={ - "required_input": _workflow.Input(_types.Types.Integer), - "default_input": _workflow.Input(_types.Types.Integer, default=5), - }, - ) - lp = workflow_to_test.create_launch_plan( - notifications=[_notifications.PagerDuty([_execution.WorkflowExecutionPhase.FAILED], ["me@myplace.com"])] - ) - assert len(lp.entity_metadata.notifications) == 1 - assert lp.entity_metadata.notifications[0].pager_duty.recipients_email == ["me@myplace.com"] - assert lp.entity_metadata.notifications[0].phases == [_execution.WorkflowExecutionPhase.FAILED] - - -def test_no_notifications(): - workflow_to_test = _workflow.workflow( - {}, - inputs={ - "required_input": _workflow.Input(_types.Types.Integer), - "default_input": _workflow.Input(_types.Types.Integer, default=5), - }, - ) - lp = workflow_to_test.create_launch_plan() - assert len(lp.entity_metadata.notifications) == 0 - - -def test_launch_plan_node(): - workflow_to_test = _workflow.workflow( - {}, - inputs={ - "required_input": _workflow.Input(_types.Types.Integer), - "default_input": _workflow.Input(_types.Types.Integer, default=5), - }, - outputs={"out": _workflow.Output([1, 2, 3], sdk_type=[_types.Types.Integer])}, - ) - lp = workflow_to_test.create_launch_plan() - - # Test that required input isn't set - with _pytest.raises(_user_exceptions.FlyteAssertion): - lp() - - # Test that positional args are rejected - with _pytest.raises(_user_exceptions.FlyteAssertion): - lp(1, 2) - - # Test that type checking works - with _pytest.raises(_user_exceptions.FlyteTypeException): - lp(required_input="abc", default_input=1) - - # Test that bad arg name is detected - with _pytest.raises(_user_exceptions.FlyteAssertion): - lp(required_input=1, bad_arg=1) - - # Test default input is accounted for - n = lp(required_input=10) - assert n.inputs[0].var == "default_input" - assert n.inputs[0].binding.scalar.primitive.integer == 5 - assert n.inputs[1].var == "required_input" - assert n.inputs[1].binding.scalar.primitive.integer == 10 - - # Test default input is overridden - n = lp(required_input=10, default_input=50) - assert n.inputs[0].var == "default_input" - assert n.inputs[0].binding.scalar.primitive.integer == 50 - assert n.inputs[1].var == "required_input" - assert n.inputs[1].binding.scalar.primitive.integer == 10 - - # Test that launch plan ID ref is flexible - lp._id = "fake" - assert n.workflow_node.launchplan_ref == "fake" - lp._id = None - - # Test that outputs are promised - n.assign_id_and_return("node-id") - assert n.outputs["out"].sdk_type.to_flyte_literal_type().collection_type.simple == _type_models.SimpleType.INTEGER - assert n.outputs["out"].var == "out" - assert n.outputs["out"].node_id == "node-id" - - -def test_labels(): - workflow_to_test = _workflow.workflow( - {}, - inputs={ - "required_input": _workflow.Input(_types.Types.Integer), - "default_input": _workflow.Input(_types.Types.Integer, default=5), - }, - ) - lp = workflow_to_test.create_launch_plan( - fixed_inputs={"required_input": 5}, - schedule=_schedules.CronSchedule("* * ? * * *"), - role="what", - labels=_common_models.Labels({"my": "label"}), - ) - assert lp.labels.values == {"my": "label"} - - -def test_annotations(): - workflow_to_test = _workflow.workflow( - {}, - inputs={ - "required_input": _workflow.Input(_types.Types.Integer), - "default_input": _workflow.Input(_types.Types.Integer, default=5), - }, - ) - lp = workflow_to_test.create_launch_plan( - fixed_inputs={"required_input": 5}, - schedule=_schedules.CronSchedule("* * ? * * *"), - role="what", - annotations=_common_models.Annotations({"my": "annotation"}), - ) - assert lp.annotations.values == {"my": "annotation"} - - -def test_serialize(): - workflow_to_test = _workflow.workflow( - {}, - inputs={ - "required_input": _workflow.Input(_types.Types.Integer), - "default_input": _workflow.Input(_types.Types.Integer, default=5), - }, - ) - workflow_to_test.id = _identifier.Identifier(_identifier.ResourceType.WORKFLOW, "p", "d", "n", "v") - lp = workflow_to_test.create_launch_plan( - fixed_inputs={"required_input": 5}, - role="iam_role", - ) - - with _configuration.TemporaryConfiguration( - _os.path.join( - _os.path.dirname(_os.path.realpath(__file__)), - "../../common/configs/local.config", - ), - internal_overrides={"image": "myflyteimage:v123", "project": "myflyteproject", "domain": "development"}, - ): - s = lp.serialize() - - assert ( - s.spec.workflow_id - == _identifier.Identifier(_identifier.ResourceType.WORKFLOW, "p", "d", "n", "v").to_flyte_idl() - ) - assert s.spec.auth_role.assumable_iam_role == "iam_role" - assert s.spec.default_inputs.parameters["default_input"].default.scalar.primitive.integer == 5 - - -def test_promote_from_model(): - workflow_to_test = _workflow.workflow( - {}, - inputs={ - "required_input": _workflow.Input(_types.Types.Integer), - "default_input": _workflow.Input(_types.Types.Integer, default=5), - }, - ) - workflow_to_test.id = _identifier.Identifier(_identifier.ResourceType.WORKFLOW, "p", "d", "n", "v") - lp = workflow_to_test.create_launch_plan( - fixed_inputs={"required_input": 5}, - schedule=_schedules.CronSchedule("* * ? * * *"), - role="what", - labels=_common_models.Labels({"my": "label"}), - ) - - with _pytest.raises(_user_exceptions.FlyteAssertion): - _launch_plan.SdkRunnableLaunchPlan.from_flyte_idl(lp.to_flyte_idl()) - - lp_from_spec = _launch_plan.SdkLaunchPlan.from_flyte_idl(lp.to_flyte_idl()) - assert not isinstance(lp_from_spec, _launch_plan.SdkRunnableLaunchPlan) - assert isinstance(lp_from_spec, _launch_plan.SdkLaunchPlan) - assert lp_from_spec == lp - - -def test_raw_data_output_prefix(): - workflow_to_test = _workflow.workflow( - {}, - inputs={ - "required_input": _workflow.Input(_types.Types.Integer), - "default_input": _workflow.Input(_types.Types.Integer, default=5), - }, - ) - lp = workflow_to_test.create_launch_plan( - fixed_inputs={"required_input": 5}, - raw_output_data_prefix="s3://bucket-name", - ) - assert lp.raw_output_data_config.output_location_prefix == "s3://bucket-name" - - lp2 = workflow_to_test.create_launch_plan( - fixed_inputs={"required_input": 5}, - ) - assert lp2.raw_output_data_config.output_location_prefix == "" diff --git a/tests/flytekit/unit/common_tests/test_nodes.py b/tests/flytekit/unit/common_tests/test_nodes.py deleted file mode 100644 index 0c462c552c..0000000000 --- a/tests/flytekit/unit/common_tests/test_nodes.py +++ /dev/null @@ -1,292 +0,0 @@ -import datetime as _datetime - -import pytest as _pytest - -from flytekit.common import component_nodes as _component_nodes -from flytekit.common import interface as _interface -from flytekit.common import nodes as _nodes -from flytekit.common.exceptions import system as _system_exceptions -from flytekit.models import literals as _literals -from flytekit.models.core import identifier as _identifier -from flytekit.models.core import workflow as _core_workflow_models -from flytekit.sdk import tasks as _tasks -from flytekit.sdk import types as _types -from flytekit.sdk import workflow as _workflow - - -def test_sdk_node_from_task(): - @_tasks.inputs(a=_types.Types.Integer) - @_tasks.outputs(b=_types.Types.Integer) - @_tasks.python_task() - def testy_test(wf_params, a, b): - pass - - n = _nodes.SdkNode( - "n", - [], - [ - _literals.Binding( - "a", - _interface.BindingData.from_python_std(_types.Types.Integer.to_flyte_literal_type(), 3), - ) - ], - _core_workflow_models.NodeMetadata("abc", _datetime.timedelta(minutes=15), _literals.RetryStrategy(3)), - sdk_task=testy_test, - sdk_workflow=None, - sdk_launch_plan=None, - sdk_branch=None, - ) - - assert n.id == "n" - assert len(n.inputs) == 1 - assert n.inputs[0].var == "a" - assert n.inputs[0].binding.scalar.primitive.integer == 3 - assert len(n.outputs) == 1 - assert "b" in n.outputs - assert n.outputs["b"].node_id == "n" - assert n.outputs["b"].var == "b" - assert n.outputs["b"].sdk_node == n - assert n.outputs["b"].sdk_type == _types.Types.Integer - assert n.metadata.name == "abc" - assert n.metadata.retries.retries == 3 - assert n.metadata.interruptible is None - assert len(n.upstream_nodes) == 0 - assert len(n.upstream_node_ids) == 0 - assert len(n.output_aliases) == 0 - - n2 = _nodes.SdkNode( - "n2", - [n], - [ - _literals.Binding( - "a", - _interface.BindingData.from_python_std(_types.Types.Integer.to_flyte_literal_type(), n.outputs.b), - ) - ], - _core_workflow_models.NodeMetadata("abc2", _datetime.timedelta(minutes=15), _literals.RetryStrategy(3)), - sdk_task=testy_test, - sdk_workflow=None, - sdk_launch_plan=None, - sdk_branch=None, - ) - - assert n2.id == "n2" - assert len(n2.inputs) == 1 - assert n2.inputs[0].var == "a" - assert n2.inputs[0].binding.promise.var == "b" - assert n2.inputs[0].binding.promise.node_id == "n" - assert len(n2.outputs) == 1 - assert "b" in n2.outputs - assert n2.outputs["b"].node_id == "n2" - assert n2.outputs["b"].var == "b" - assert n2.outputs["b"].sdk_node == n2 - assert n2.outputs["b"].sdk_type == _types.Types.Integer - assert n2.metadata.name == "abc2" - assert n2.metadata.retries.retries == 3 - assert "n" in n2.upstream_node_ids - assert n in n2.upstream_nodes - assert len(n2.upstream_nodes) == 1 - assert len(n2.upstream_node_ids) == 1 - assert len(n2.output_aliases) == 0 - - # Test right shift operator and late binding - n3 = _nodes.SdkNode( - "n3", - [], - [ - _literals.Binding( - "a", - _interface.BindingData.from_python_std(_types.Types.Integer.to_flyte_literal_type(), 3), - ) - ], - _core_workflow_models.NodeMetadata("abc3", _datetime.timedelta(minutes=15), _literals.RetryStrategy(3)), - sdk_task=testy_test, - sdk_workflow=None, - sdk_launch_plan=None, - sdk_branch=None, - ) - n2 >> n3 - n >> n2 >> n3 - n3 << n2 - n3 << n2 << n - - assert n3.id == "n3" - assert len(n3.inputs) == 1 - assert n3.inputs[0].var == "a" - assert n3.inputs[0].binding.scalar.primitive.integer == 3 - assert len(n3.outputs) == 1 - assert "b" in n3.outputs - assert n3.outputs["b"].node_id == "n3" - assert n3.outputs["b"].var == "b" - assert n3.outputs["b"].sdk_node == n3 - assert n3.outputs["b"].sdk_type == _types.Types.Integer - assert n3.metadata.name == "abc3" - assert n3.metadata.retries.retries == 3 - assert "n2" in n3.upstream_node_ids - assert n2 in n3.upstream_nodes - assert len(n3.upstream_nodes) == 1 - assert len(n3.upstream_node_ids) == 1 - assert len(n3.output_aliases) == 0 - - # Test left shift operator and late binding - n4 = _nodes.SdkNode( - "n4", - [], - [ - _literals.Binding( - "a", - _interface.BindingData.from_python_std(_types.Types.Integer.to_flyte_literal_type(), 3), - ) - ], - _core_workflow_models.NodeMetadata("abc4", _datetime.timedelta(minutes=15), _literals.RetryStrategy(3)), - sdk_task=testy_test, - sdk_workflow=None, - sdk_launch_plan=None, - sdk_branch=None, - ) - - n4 << n3 - - # Test that implicit dependencies don't cause direct dependencies - n4 << n3 << n2 << n - n >> n2 >> n3 >> n4 - - assert n4.id == "n4" - assert len(n4.inputs) == 1 - assert n4.inputs[0].var == "a" - assert n4.inputs[0].binding.scalar.primitive.integer == 3 - assert len(n4.outputs) == 1 - assert "b" in n4.outputs - assert n4.outputs["b"].node_id == "n4" - assert n4.outputs["b"].var == "b" - assert n4.outputs["b"].sdk_node == n4 - assert n4.outputs["b"].sdk_type == _types.Types.Integer - assert n4.metadata.name == "abc4" - assert n4.metadata.retries.retries == 3 - assert "n3" in n4.upstream_node_ids - assert n3 in n4.upstream_nodes - assert len(n4.upstream_nodes) == 1 - assert len(n4.upstream_node_ids) == 1 - assert len(n4.output_aliases) == 0 - - # Add another dependency - n4 << n2 - assert "n3" in n4.upstream_node_ids - assert n3 in n4.upstream_nodes - assert "n2" in n4.upstream_node_ids - assert n2 in n4.upstream_nodes - assert len(n4.upstream_nodes) == 2 - assert len(n4.upstream_node_ids) == 2 - - -def test_sdk_task_node(): - @_tasks.inputs(a=_types.Types.Integer) - @_tasks.outputs(b=_types.Types.Integer) - @_tasks.python_task() - def testy_test(wf_params, a, b): - pass - - testy_test._id = _identifier.Identifier(_identifier.ResourceType.TASK, "project", "domain", "name", "version") - n = _component_nodes.SdkTaskNode(testy_test) - assert n.reference_id.project == "project" - assert n.reference_id.domain == "domain" - assert n.reference_id.name == "name" - assert n.reference_id.version == "version" - - # Test floating ID - testy_test._id = _identifier.Identifier( - _identifier.ResourceType.TASK, - "new_project", - "new_domain", - "new_name", - "new_version", - ) - assert n.reference_id.project == "new_project" - assert n.reference_id.domain == "new_domain" - assert n.reference_id.name == "new_name" - assert n.reference_id.version == "new_version" - - -def test_sdk_node_from_lp(): - @_tasks.inputs(a=_types.Types.Integer) - @_tasks.outputs(b=_types.Types.Integer) - @_tasks.python_task() - def testy_test(wf_params, a, b): - pass - - @_workflow.workflow_class - class test_workflow(object): - a = _workflow.Input(_types.Types.Integer) - test = testy_test(a=a) - b = _workflow.Output(test.outputs.b, sdk_type=_types.Types.Integer) - - lp = test_workflow.create_launch_plan() - - n1 = _nodes.SdkNode( - "n1", - [], - [ - _literals.Binding( - "a", - _interface.BindingData.from_python_std(_types.Types.Integer.to_flyte_literal_type(), 3), - ) - ], - _core_workflow_models.NodeMetadata("abc", _datetime.timedelta(minutes=15), _literals.RetryStrategy(3)), - sdk_launch_plan=lp, - ) - - assert n1.id == "n1" - assert len(n1.inputs) == 1 - assert n1.inputs[0].var == "a" - assert n1.inputs[0].binding.scalar.primitive.integer == 3 - assert len(n1.outputs) == 1 - assert "b" in n1.outputs - assert n1.outputs["b"].node_id == "n1" - assert n1.outputs["b"].var == "b" - assert n1.outputs["b"].sdk_node == n1 - assert n1.outputs["b"].sdk_type == _types.Types.Integer - assert n1.metadata.name == "abc" - assert n1.metadata.retries.retries == 3 - assert len(n1.upstream_nodes) == 0 - assert len(n1.upstream_node_ids) == 0 - assert len(n1.output_aliases) == 0 - - -def test_sdk_launch_plan_node(): - @_tasks.inputs(a=_types.Types.Integer) - @_tasks.outputs(b=_types.Types.Integer) - @_tasks.python_task() - def testy_test(wf_params, a, b): - pass - - @_workflow.workflow_class - class test_workflow(object): - a = _workflow.Input(_types.Types.Integer) - test = testy_test(a=1) - b = _workflow.Output(test.outputs.b, sdk_type=_types.Types.Integer) - - lp = test_workflow.create_launch_plan() - - lp._id = _identifier.Identifier(_identifier.ResourceType.TASK, "project", "domain", "name", "version") - n = _component_nodes.SdkWorkflowNode(sdk_launch_plan=lp) - assert n.launchplan_ref.project == "project" - assert n.launchplan_ref.domain == "domain" - assert n.launchplan_ref.name == "name" - assert n.launchplan_ref.version == "version" - - # Test floating ID - lp._id = _identifier.Identifier( - _identifier.ResourceType.TASK, - "new_project", - "new_domain", - "new_name", - "new_version", - ) - assert n.launchplan_ref.project == "new_project" - assert n.launchplan_ref.domain == "new_domain" - assert n.launchplan_ref.name == "new_name" - assert n.launchplan_ref.version == "new_version" - - # If you specify both, you should get an exception - with _pytest.raises(_system_exceptions.FlyteSystemException): - _component_nodes.SdkWorkflowNode(sdk_workflow=test_workflow, sdk_launch_plan=lp) diff --git a/tests/flytekit/unit/common_tests/test_notifications.py b/tests/flytekit/unit/common_tests/test_notifications.py deleted file mode 100644 index 8a2278097a..0000000000 --- a/tests/flytekit/unit/common_tests/test_notifications.py +++ /dev/null @@ -1,35 +0,0 @@ -from flytekit.common import notifications as _notifications -from flytekit.models.core import execution as _execution_model - - -def test_pager_duty(): - obj = _notifications.PagerDuty([_execution_model.WorkflowExecutionPhase.FAILED], ["me@myplace.com"]) - assert obj.email is None - assert obj.slack is None - assert obj.phases == [_execution_model.WorkflowExecutionPhase.FAILED] - assert obj.pager_duty.recipients_email == ["me@myplace.com"] - - obj2 = _notifications.PagerDuty.from_flyte_idl(obj.to_flyte_idl()) - assert obj == obj2 - - -def test_slack(): - obj = _notifications.Slack([_execution_model.WorkflowExecutionPhase.FAILED], ["me@myplace.com"]) - assert obj.email is None - assert obj.pager_duty is None - assert obj.phases == [_execution_model.WorkflowExecutionPhase.FAILED] - assert obj.slack.recipients_email == ["me@myplace.com"] - - obj2 = _notifications.Slack.from_flyte_idl(obj.to_flyte_idl()) - assert obj == obj2 - - -def test_email(): - obj = _notifications.Email([_execution_model.WorkflowExecutionPhase.FAILED], ["me@myplace.com"]) - assert obj.pager_duty is None - assert obj.slack is None - assert obj.phases == [_execution_model.WorkflowExecutionPhase.FAILED] - assert obj.email.recipients_email == ["me@myplace.com"] - - obj2 = _notifications.Email.from_flyte_idl(obj.to_flyte_idl()) - assert obj == obj2 diff --git a/tests/flytekit/unit/common_tests/test_promise.py b/tests/flytekit/unit/common_tests/test_promise.py deleted file mode 100644 index e06cca7dfa..0000000000 --- a/tests/flytekit/unit/common_tests/test_promise.py +++ /dev/null @@ -1,89 +0,0 @@ -import pytest - -from flytekit import FlyteContextManager -from flytekit.common import promise -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.types import base_sdk_types, primitives -from flytekit.core.interface import Interface -from flytekit.core.promise import Promise, create_native_named_tuple, extract_obj_name -from flytekit.core.type_engine import TypeEngine -from flytekit.models.types import LiteralType, SimpleType - - -def test_input(): - i = promise.Input("name", primitives.Integer, help="blah", default=None) - assert i.name == "name" - assert i.sdk_default is None - assert i.default == base_sdk_types.Void() - assert i.sdk_required is False - assert i.help == "blah" - assert i.var.description == "blah" - assert i.sdk_type == primitives.Integer - - i = promise.Input("name2", primitives.Integer, default=1) - assert i.name == "name2" - assert i.sdk_default == 1 - assert i.default == primitives.Integer(1) - assert i.required is None - assert i.sdk_required is False - assert i.help is None - assert i.var.description == "" - assert i.sdk_type == primitives.Integer - - with pytest.raises(_user_exceptions.FlyteAssertion): - promise.Input("abc", primitives.Integer, required=True, default=1) - - -def test_create_native_named_tuple(): - ctx = FlyteContextManager.current_context() - t = create_native_named_tuple(ctx, promises=None, entity_interface=Interface()) - assert t is None - - p1 = Promise(var="x", val=TypeEngine.to_literal(ctx, 1, int, LiteralType(simple=SimpleType.INTEGER))) - p2 = Promise(var="y", val=TypeEngine.to_literal(ctx, 2, int, LiteralType(simple=SimpleType.INTEGER))) - - t = create_native_named_tuple(ctx, promises=p1, entity_interface=Interface(outputs={"x": int})) - assert t - assert t == 1 - - t = create_native_named_tuple(ctx, promises=[], entity_interface=Interface()) - assert t is None - - t = create_native_named_tuple(ctx, promises=[p1, p2], entity_interface=Interface(outputs={"x": int, "y": int})) - assert t - assert t == (1, 2) - - t = create_native_named_tuple( - ctx, promises=[p1, p2], entity_interface=Interface(outputs={"x": int, "y": int}, output_tuple_name="Tup") - ) - assert t - assert t == (1, 2) - assert t.__class__.__name__ == "Tup" - - with pytest.raises(KeyError): - create_native_named_tuple( - ctx, promises=[p1, p2], entity_interface=Interface(outputs={"x": int}, output_tuple_name="Tup") - ) - - with pytest.raises(AssertionError, match="Failed to convert value of output x"): - create_native_named_tuple(ctx, promises=[p1, p2], entity_interface=Interface(outputs={"x": Promise, "y": int})) - - with pytest.raises(AssertionError, match="Failed to convert value of output x"): - create_native_named_tuple(ctx, promises=p1, entity_interface=Interface(outputs={"x": Promise})) - - -@pytest.mark.parametrize( - "name, expected_name", - [ - ("test", "test"), - ("test.abc", "abc"), - (".test", "test"), - ("test.", ""), - ("test.xyz.abc", "abc"), - ("", ""), - (None, ""), - ("test.xyz.abc.", ""), - ], -) -def test_extract_obj_name(name, expected_name): - assert extract_obj_name(name) == expected_name diff --git a/tests/flytekit/unit/common_tests/test_schedules.py b/tests/flytekit/unit/common_tests/test_schedules.py deleted file mode 100644 index 4e6f231307..0000000000 --- a/tests/flytekit/unit/common_tests/test_schedules.py +++ /dev/null @@ -1,131 +0,0 @@ -import datetime as _datetime - -import pytest as _pytest - -from flytekit.common import schedules as _schedules -from flytekit.common.exceptions import user as _user_exceptions - - -def test_cron(): - obj = _schedules.CronSchedule("* * ? * * *", kickoff_time_input_arg="abc") - assert obj.kickoff_time_input_arg == "abc" - assert obj.cron_expression == "* * ? * * *" - assert obj == _schedules.CronSchedule.from_flyte_idl(obj.to_flyte_idl()) - - -def test_cron_karg(): - obj = _schedules.CronSchedule(cron_expression="* * ? * * *", kickoff_time_input_arg="abc") - assert obj.kickoff_time_input_arg == "abc" - assert obj.cron_expression == "* * ? * * *" - assert obj == _schedules.CronSchedule.from_flyte_idl(obj.to_flyte_idl()) - - -def test_cron_validation(): - with _pytest.raises(_user_exceptions.FlyteAssertion): - _schedules.CronSchedule("* * * * * *", kickoff_time_input_arg="abc") - - with _pytest.raises(_user_exceptions.FlyteAssertion): - _schedules.CronSchedule("* * ? * *", kickoff_time_input_arg="abc") - - -def test_fixed_rate(): - obj = _schedules.FixedRate(_datetime.timedelta(hours=10), kickoff_time_input_arg="abc") - assert obj.rate.unit == _schedules.FixedRate.FixedRateUnit.HOUR - assert obj.rate.value == 10 - assert obj == _schedules.FixedRate.from_flyte_idl(obj.to_flyte_idl()) - - obj = _schedules.FixedRate(_datetime.timedelta(hours=24), kickoff_time_input_arg="abc") - assert obj.rate.unit == _schedules.FixedRate.FixedRateUnit.DAY - assert obj.rate.value == 1 - assert obj == _schedules.FixedRate.from_flyte_idl(obj.to_flyte_idl()) - - obj = _schedules.FixedRate(_datetime.timedelta(minutes=30), kickoff_time_input_arg="abc") - assert obj.rate.unit == _schedules.FixedRate.FixedRateUnit.MINUTE - assert obj.rate.value == 30 - assert obj == _schedules.FixedRate.from_flyte_idl(obj.to_flyte_idl()) - - obj = _schedules.FixedRate(_datetime.timedelta(minutes=120), kickoff_time_input_arg="abc") - assert obj.rate.unit == _schedules.FixedRate.FixedRateUnit.HOUR - assert obj.rate.value == 2 - assert obj == _schedules.FixedRate.from_flyte_idl(obj.to_flyte_idl()) - - -def test_fixed_rate_bad_duration(): - pass - - -def test_fixed_rate_negative_duration(): - pass - - -@_pytest.mark.parametrize( - "schedule", - [ - "hourly", - "hours", - "HOURS", - "@hourly", - "daily", - "days", - "DAYS", - "@daily", - "weekly", - "weeks", - "WEEKS", - "@weekly", - "monthly", - "months", - "MONTHS", - "@monthly", - "annually", - "@annually", - "yearly", - "years", - "YEARS", - "@yearly", - "* * * * *", - ], -) -def test_cron_schedule_schedule_validation(schedule): - obj = _schedules.CronSchedule(schedule=schedule, kickoff_time_input_arg="abc") - assert obj.cron_schedule.schedule == schedule - - -@_pytest.mark.parametrize( - "schedule", - ["foo", "* *"], -) -def test_cron_schedule_schedule_validation_invalid(schedule): - with _pytest.raises(_user_exceptions.FlyteAssertion): - _schedules.CronSchedule(schedule=schedule, kickoff_time_input_arg="abc") - - -def test_cron_schedule_offset_validation_invalid(): - with _pytest.raises(_user_exceptions.FlyteAssertion): - _schedules.CronSchedule(schedule="days", offset="foo", kickoff_time_input_arg="abc") - - -def test_cron_schedule(): - obj = _schedules.CronSchedule(schedule="days", kickoff_time_input_arg="abc") - assert obj.cron_schedule.schedule == "days" - assert obj.cron_schedule.offset is None - assert obj == _schedules.CronSchedule.from_flyte_idl(obj.to_flyte_idl()) - - -def test_cron_schedule_offset(): - obj = _schedules.CronSchedule(schedule="days", offset="P1D", kickoff_time_input_arg="abc") - assert obj.cron_schedule.schedule == "days" - assert obj.cron_schedule.offset == "P1D" - assert obj == _schedules.CronSchedule.from_flyte_idl(obj.to_flyte_idl()) - - -def test_both_cron_expression_and_cron_schedule_schedule(): - with _pytest.raises(_user_exceptions.FlyteAssertion): - _schedules.CronSchedule( - cron_expression="* * ? * * *", schedule="days", offset="foo", kickoff_time_input_arg="abc" - ) - - -def test_cron_expression_and_cron_schedule_offset(): - with _pytest.raises(_user_exceptions.FlyteAssertion): - _schedules.CronSchedule(cron_expression="* * ? * * *", offset="foo", kickoff_time_input_arg="abc") diff --git a/tests/flytekit/unit/common_tests/test_workflow.py b/tests/flytekit/unit/common_tests/test_workflow.py deleted file mode 100644 index 13a9e65d8e..0000000000 --- a/tests/flytekit/unit/common_tests/test_workflow.py +++ /dev/null @@ -1,383 +0,0 @@ -import pytest as _pytest -from flyteidl.admin import workflow_pb2 as _workflow_pb2 - -from flytekit.common import constants, interface -from flytekit.common import local_workflow as _local_workflow -from flytekit.common import nodes, promise, workflow -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.local_workflow import build_sdk_workflow_from_metaclass -from flytekit.common.types import containers, primitives -from flytekit.models import literals as _literals -from flytekit.models.core import identifier as _identifier -from flytekit.models.core import workflow as _workflow_models -from flytekit.sdk import types as _types -from flytekit.sdk.tasks import inputs, outputs, python_task - - -def test_output(): - o = _local_workflow.Output("name", 1, sdk_type=primitives.Integer, help="blah") - assert o.name == "name" - assert o.var.description == "blah" - assert o.var.type == primitives.Integer.to_flyte_literal_type() - assert o.binding_data.scalar.primitive.integer == 1 - - -def test_workflow(): - @inputs(a=primitives.Integer) - @outputs(b=primitives.Integer) - @python_task() - def my_task(wf_params, a, b): - b.set(a + 1) - - my_task._id = _identifier.Identifier(_identifier.ResourceType.TASK, "project", "domain", "my_task", "version") - - @inputs(a=[primitives.Integer]) - @outputs(b=[primitives.Integer]) - @python_task - def my_list_task(wf_params, a, b): - b.set([v + 1 for v in a]) - - my_list_task._id = _identifier.Identifier( - _identifier.ResourceType.TASK, "project", "domain", "my_list_task", "version" - ) - - input_list = [ - promise.Input("input_1", primitives.Integer), - promise.Input("input_2", primitives.Integer, default=5, help="Not required."), - ] - - n1 = my_task(a=input_list[0]).assign_id_and_return("n1") - n2 = my_task(a=input_list[1]).assign_id_and_return("n2") - n3 = my_task(a=100).assign_id_and_return("n3") - n4 = my_task(a=n1.outputs.b).assign_id_and_return("n4") - n5 = my_list_task(a=[input_list[0], input_list[1], n3.outputs.b, 100]).assign_id_and_return("n5") - n6 = my_list_task(a=n5.outputs.b) - n1 >> n6 - - nodes = [n1, n2, n3, n4, n5, n6] - - w = _local_workflow.SdkRunnableWorkflow.construct_from_class_definition( - inputs=input_list, - outputs=[_local_workflow.Output("a", n1.outputs.b, sdk_type=primitives.Integer)], - nodes=nodes, - ) - - assert w.interface.inputs["input_1"].type == primitives.Integer.to_flyte_literal_type() - assert w.interface.inputs["input_2"].type == primitives.Integer.to_flyte_literal_type() - assert w.nodes[0].inputs[0].var == "a" - assert w.nodes[0].inputs[0].binding.promise.node_id == constants.GLOBAL_INPUT_NODE_ID - assert w.nodes[0].inputs[0].binding.promise.var == "input_1" - assert w.nodes[1].inputs[0].binding.promise.node_id == constants.GLOBAL_INPUT_NODE_ID - assert w.nodes[1].inputs[0].binding.promise.var == "input_2" - assert w.nodes[2].inputs[0].binding.scalar.primitive.integer == 100 - assert w.nodes[3].inputs[0].var == "a" - assert w.nodes[3].inputs[0].binding.promise.node_id == n1.id - - # Test conversion to flyte_idl and back - w._id = _identifier.Identifier(_identifier.ResourceType.WORKFLOW, "fake", "faker", "fakest", "fakerest") - w = _workflow_models.WorkflowTemplate.from_flyte_idl(w.to_flyte_idl()) - assert w.interface.inputs["input_1"].type == primitives.Integer.to_flyte_literal_type() - assert w.interface.inputs["input_2"].type == primitives.Integer.to_flyte_literal_type() - assert w.nodes[0].inputs[0].var == "a" - assert w.nodes[0].inputs[0].binding.promise.node_id == constants.GLOBAL_INPUT_NODE_ID - assert w.nodes[0].inputs[0].binding.promise.var == "input_1" - assert w.nodes[1].inputs[0].binding.promise.node_id == constants.GLOBAL_INPUT_NODE_ID - assert w.nodes[1].inputs[0].binding.promise.var == "input_2" - assert w.nodes[2].inputs[0].binding.scalar.primitive.integer == 100 - assert w.nodes[3].inputs[0].var == "a" - assert w.nodes[3].inputs[0].binding.promise.node_id == n1.id - assert w.nodes[4].inputs[0].var == "a" - assert w.nodes[4].inputs[0].binding.collection.bindings[0].promise.node_id == constants.GLOBAL_INPUT_NODE_ID - assert w.nodes[4].inputs[0].binding.collection.bindings[0].promise.var == "input_1" - assert w.nodes[4].inputs[0].binding.collection.bindings[1].promise.node_id == constants.GLOBAL_INPUT_NODE_ID - assert w.nodes[4].inputs[0].binding.collection.bindings[1].promise.var == "input_2" - assert w.nodes[4].inputs[0].binding.collection.bindings[2].promise.node_id == n3.id - assert w.nodes[4].inputs[0].binding.collection.bindings[2].promise.var == "b" - assert w.nodes[4].inputs[0].binding.collection.bindings[3].scalar.primitive.integer == 100 - assert w.nodes[5].inputs[0].var == "a" - assert w.nodes[5].inputs[0].binding.promise.node_id == n5.id - assert w.nodes[5].inputs[0].binding.promise.var == "b" - - assert len(w.outputs) == 1 - assert w.outputs[0].var == "a" - assert w.outputs[0].binding.promise.var == "b" - assert w.outputs[0].binding.promise.node_id == "n1" - # TODO: Test promotion of w -> SdkWorkflow - - -def test_workflow_decorator(): - @inputs(a=primitives.Integer) - @outputs(b=primitives.Integer) - @python_task - def my_task(wf_params, a, b): - b.set(a + 1) - - my_task._id = _identifier.Identifier(_identifier.ResourceType.TASK, "propject", "domain", "my_task", "version") - - @inputs(a=[primitives.Integer]) - @outputs(b=[primitives.Integer]) - @python_task - def my_list_task(wf_params, a, b): - b.set([v + 1 for v in a]) - - my_list_task._id = _identifier.Identifier( - _identifier.ResourceType.TASK, "propject", "domain", "my_list_task", "version" - ) - - class my_workflow(object): - input_1 = promise.Input("input_1", primitives.Integer) - input_2 = promise.Input("input_2", primitives.Integer, default=5, help="Not required.") - n1 = my_task(a=input_1) - n2 = my_task(a=input_2) - n3 = my_task(a=100) - n4 = my_task(a=n1.outputs.b) - n5 = my_list_task(a=[input_1, input_2, n3.outputs.b, 100]) - n6 = my_list_task(a=n5.outputs.b) - n1 >> n6 - a = _local_workflow.Output("a", n1.outputs.b, sdk_type=primitives.Integer) - - w = _local_workflow.build_sdk_workflow_from_metaclass( - my_workflow, - on_failure=_workflow_models.WorkflowMetadata.OnFailurePolicy.FAIL_AFTER_EXECUTABLE_NODES_COMPLETE, - ) - - assert w.should_create_default_launch_plan is True - - assert w.interface.inputs["input_1"].type == primitives.Integer.to_flyte_literal_type() - assert w.interface.inputs["input_2"].type == primitives.Integer.to_flyte_literal_type() - assert w.nodes[0].inputs[0].var == "a" - assert w.nodes[0].inputs[0].binding.promise.node_id == constants.GLOBAL_INPUT_NODE_ID - assert w.nodes[0].inputs[0].binding.promise.var == "input_1" - assert w.nodes[1].inputs[0].binding.promise.node_id == constants.GLOBAL_INPUT_NODE_ID - assert w.nodes[1].inputs[0].binding.promise.var == "input_2" - assert w.nodes[2].inputs[0].binding.scalar.primitive.integer == 100 - assert w.nodes[3].inputs[0].var == "a" - assert w.nodes[3].inputs[0].binding.promise.node_id == "n1" - - # Test conversion to flyte_idl and back - w.id = _identifier.Identifier(_identifier.ResourceType.WORKFLOW, "fake", "faker", "fakest", "fakerest") - w = _workflow_models.WorkflowTemplate.from_flyte_idl(w.to_flyte_idl()) - assert w.interface.inputs["input_1"].type == primitives.Integer.to_flyte_literal_type() - assert w.interface.inputs["input_2"].type == primitives.Integer.to_flyte_literal_type() - assert w.nodes[0].inputs[0].var == "a" - assert w.nodes[0].inputs[0].binding.promise.node_id == constants.GLOBAL_INPUT_NODE_ID - assert w.nodes[0].inputs[0].binding.promise.var == "input_1" - assert w.nodes[1].inputs[0].binding.promise.node_id == constants.GLOBAL_INPUT_NODE_ID - assert w.nodes[1].inputs[0].binding.promise.var == "input_2" - assert w.nodes[2].inputs[0].binding.scalar.primitive.integer == 100 - assert w.nodes[3].inputs[0].var == "a" - assert w.nodes[3].inputs[0].binding.promise.node_id == "n1" - assert w.nodes[4].inputs[0].var == "a" - assert w.nodes[4].inputs[0].binding.collection.bindings[0].promise.node_id == constants.GLOBAL_INPUT_NODE_ID - assert w.nodes[4].inputs[0].binding.collection.bindings[0].promise.var == "input_1" - assert w.nodes[4].inputs[0].binding.collection.bindings[1].promise.node_id == constants.GLOBAL_INPUT_NODE_ID - assert w.nodes[4].inputs[0].binding.collection.bindings[1].promise.var == "input_2" - assert w.nodes[4].inputs[0].binding.collection.bindings[2].promise.node_id == "n3" - assert w.nodes[4].inputs[0].binding.collection.bindings[2].promise.var == "b" - assert w.nodes[4].inputs[0].binding.collection.bindings[3].scalar.primitive.integer == 100 - assert w.nodes[5].inputs[0].var == "a" - assert w.nodes[5].inputs[0].binding.promise.node_id == "n5" - assert w.nodes[5].inputs[0].binding.promise.var == "b" - - assert len(w.outputs) == 1 - assert w.outputs[0].var == "a" - assert w.outputs[0].binding.promise.var == "b" - assert w.outputs[0].binding.promise.node_id == "n1" - assert ( - w.metadata.on_failure == _workflow_models.WorkflowMetadata.OnFailurePolicy.FAIL_AFTER_EXECUTABLE_NODES_COMPLETE - ) - # TODO: Test promotion of w -> SdkWorkflow - - -def test_workflow_node(): - @inputs(a=primitives.Integer) - @outputs(b=primitives.Integer) - @python_task() - def my_task(wf_params, a, b): - b.set(a + 1) - - my_task._id = _identifier.Identifier(_identifier.ResourceType.TASK, "project", "domain", "my_task", "version") - - @inputs(a=[primitives.Integer]) - @outputs(b=[primitives.Integer]) - @python_task - def my_list_task(wf_params, a, b): - b.set([v + 1 for v in a]) - - my_list_task._id = _identifier.Identifier( - _identifier.ResourceType.TASK, "project", "domain", "my_list_task", "version" - ) - - input_list = [ - promise.Input("required", primitives.Integer), - promise.Input("not_required", primitives.Integer, default=5, help="Not required."), - ] - - n1 = my_task(a=input_list[0]).assign_id_and_return("n1") - n2 = my_task(a=input_list[1]).assign_id_and_return("n2") - n3 = my_task(a=100).assign_id_and_return("n3") - n4 = my_task(a=n1.outputs.b).assign_id_and_return("n4") - n5 = my_list_task(a=[input_list[0], input_list[1], n3.outputs.b, 100]).assign_id_and_return("n5") - n6 = my_list_task(a=n5.outputs.b) - - nodes = [n1, n2, n3, n4, n5, n6] - - wf_out = [ - _local_workflow.Output( - "nested_out", - [n5.outputs.b, n6.outputs.b, [n1.outputs.b, n2.outputs.b]], - sdk_type=[[primitives.Integer]], - ), - _local_workflow.Output("scalar_out", n1.outputs.b, sdk_type=primitives.Integer), - ] - - w = _local_workflow.SdkRunnableWorkflow.construct_from_class_definition( - inputs=input_list, outputs=wf_out, nodes=nodes - ) - - # Test that required input isn't set - with _pytest.raises(_user_exceptions.FlyteAssertion): - w() - - # Test that positional args are rejected - with _pytest.raises(_user_exceptions.FlyteAssertion): - w(1, 2) - - # Test that type checking works - with _pytest.raises(_user_exceptions.FlyteTypeException): - w(required="abc", not_required=1) - - # Test that bad arg name is detected - with _pytest.raises(_user_exceptions.FlyteAssertion): - w(required=1, bad_arg=1) - - # Test default input is accounted for - n = w(required=10) - assert n.inputs[0].var == "not_required" - assert n.inputs[0].binding.scalar.primitive.integer == 5 - assert n.inputs[1].var == "required" - assert n.inputs[1].binding.scalar.primitive.integer == 10 - - # Test default input is overridden - n = w(required=10, not_required=50) - assert n.inputs[0].var == "not_required" - assert n.inputs[0].binding.scalar.primitive.integer == 50 - assert n.inputs[1].var == "required" - assert n.inputs[1].binding.scalar.primitive.integer == 10 - - # Test that workflow is saved in the node - w.id = "fake" - assert n.workflow_node.sub_workflow_ref == "fake" - w.id = None - - # Test that outputs are promised - n.assign_id_and_return("node-id*") # dns'ified - assert n.outputs["scalar_out"].sdk_type.to_flyte_literal_type() == primitives.Integer.to_flyte_literal_type() - assert n.outputs["scalar_out"].var == "scalar_out" - assert n.outputs["scalar_out"].node_id == "node-id" - - assert ( - n.outputs["nested_out"].sdk_type.to_flyte_literal_type() - == containers.List(containers.List(primitives.Integer)).to_flyte_literal_type() - ) - assert n.outputs["nested_out"].var == "nested_out" - assert n.outputs["nested_out"].node_id == "node-id" - - -def test_non_system_nodes(): - @inputs(a=primitives.Integer) - @outputs(b=primitives.Integer) - @python_task() - def my_task(wf_params, a, b): - b.set(a + 1) - - my_task._id = _identifier.Identifier(_identifier.ResourceType.TASK, "project", "domain", "my_task", "version") - - required_input = promise.Input("required", primitives.Integer) - - n1 = my_task(a=required_input).assign_id_and_return("n1") - - n_start = nodes.SdkNode( - "start-node", - [], - [ - _literals.Binding( - "a", - interface.BindingData.from_python_std(_types.Types.Integer.to_flyte_literal_type(), 3), - ) - ], - None, - sdk_task=my_task, - sdk_workflow=None, - sdk_launch_plan=None, - sdk_branch=None, - ) - - non_system_nodes = workflow.SdkWorkflow.get_non_system_nodes([n1, n_start]) - assert len(non_system_nodes) == 1 - assert non_system_nodes[0].id == "n1" - - -def test_workflow_serialization(): - @inputs(a=primitives.Integer) - @outputs(b=primitives.Integer) - @python_task() - def my_task(wf_params, a, b): - b.set(a + 1) - - my_task._id = _identifier.Identifier(_identifier.ResourceType.TASK, "project", "domain", "my_task", "version") - - @inputs(a=[primitives.Integer]) - @outputs(b=[primitives.Integer]) - @python_task - def my_list_task(wf_params, a, b): - b.set([v + 1 for v in a]) - - my_list_task._id = _identifier.Identifier( - _identifier.ResourceType.TASK, "project", "domain", "my_list_task", "version" - ) - - input_list = [ - promise.Input("required", primitives.Integer), - promise.Input("not_required", primitives.Integer, default=5, help="Not required."), - ] - - n1 = my_task(a=input_list[0]).assign_id_and_return("n1") - n2 = my_task(a=input_list[1]).assign_id_and_return("n2") - n3 = my_task(a=100).assign_id_and_return("n3") - n4 = my_task(a=n1.outputs.b).assign_id_and_return("n4") - n5 = my_list_task(a=[input_list[0], input_list[1], n3.outputs.b, 100]).assign_id_and_return("n5") - n6 = my_list_task(a=n5.outputs.b) - - nodes = [n1, n2, n3, n4, n5, n6] - - wf_out = [ - _local_workflow.Output( - "nested_out", - [n5.outputs.b, n6.outputs.b, [n1.outputs.b, n2.outputs.b]], - sdk_type=[[primitives.Integer]], - ), - _local_workflow.Output("scalar_out", n1.outputs.b, sdk_type=primitives.Integer), - ] - - w = _local_workflow.SdkRunnableWorkflow.construct_from_class_definition( - inputs=input_list, outputs=wf_out, nodes=nodes - ) - serialized = w.serialize() - assert isinstance(serialized, _workflow_pb2.WorkflowSpec) - assert len(serialized.template.nodes) == 6 - assert len(serialized.template.interface.inputs.variables.keys()) == 2 - assert len(serialized.template.interface.outputs.variables.keys()) == 2 - - -def test_workflow_disable_default_launch_plan(): - class MyWorkflow(object): - input_1 = promise.Input("input_1", primitives.Integer) - input_2 = promise.Input("input_2", primitives.Integer, default=5, help="Not required.") - - w = build_sdk_workflow_from_metaclass( - MyWorkflow, - disable_default_launch_plan=True, - ) - - assert w.should_create_default_launch_plan is False diff --git a/tests/flytekit/unit/common_tests/test_workflow_promote.py b/tests/flytekit/unit/common_tests/test_workflow_promote.py index f165f4c231..f10ff22f1a 100644 --- a/tests/flytekit/unit/common_tests/test_workflow_promote.py +++ b/tests/flytekit/unit/common_tests/test_workflow_promote.py @@ -3,23 +3,11 @@ from flyteidl.core import compiler_pb2 as _compiler_pb2 from flyteidl.core import workflow_pb2 as _workflow_pb2 -from mock import patch as _patch -from flytekit.common import workflow as _workflow_common -from flytekit.common.tasks import task as _task -from flytekit.models import interface as _interface from flytekit.models import literals as _literals from flytekit.models import task as _task_model -from flytekit.models import types as _types from flytekit.models.core import compiler as _compiler_model -from flytekit.models.core import identifier as _identifier from flytekit.models.core import workflow as _workflow_model -from flytekit.sdk import tasks as _sdk_tasks -from flytekit.sdk import workflow as _sdk_workflow -from flytekit.sdk.tasks import inputs, outputs, python_task -from flytekit.sdk.types import Types -from flytekit.sdk.types import Types as _Types -from flytekit.sdk.workflow import Input, Output, workflow_class def get_sample_node_metadata(node_id): @@ -108,60 +96,6 @@ class OneTaskWFForPromote(object): return wt -@_patch("flytekit.common.tasks.task.SdkTask.fetch") -def test_basic_workflow_promote(mock_task_fetch): - # This section defines a sample workflow from a user - @_sdk_tasks.inputs(a=_Types.Integer) - @_sdk_tasks.outputs(b=_Types.Integer, c=_Types.Integer) - @_sdk_tasks.python_task() - def demo_task_for_promote(wf_params, a, b, c): - b.set(a + 1) - c.set(a + 2) - - @_sdk_workflow.workflow_class() - class TestPromoteExampleWf(object): - wf_input = _sdk_workflow.Input(_Types.Integer, required=True) - my_task_node = demo_task_for_promote(a=wf_input) - wf_output_b = _sdk_workflow.Output(my_task_node.outputs.b, sdk_type=_Types.Integer) - wf_output_c = _sdk_workflow.Output(my_task_node.outputs.c, sdk_type=_Types.Integer) - - # This section uses the TaskTemplate stored in Admin to promote back to an Sdk Workflow - int_type = _types.LiteralType(_types.SimpleType.INTEGER) - task_interface = _interface.TypedInterface( - # inputs - {"a": _interface.Variable(int_type, "description1")}, - # outputs - {"b": _interface.Variable(int_type, "description2"), "c": _interface.Variable(int_type, "description3")}, - ) - # Since the promotion of a workflow requires retrieving the task from Admin, we mock the SdkTask to return - task_template = _task_model.TaskTemplate( - _identifier.Identifier( - _identifier.ResourceType.TASK, - "project", - "domain", - "tests.flytekit.unit.common_tests.test_workflow_promote.demo_task_for_promote", - "version", - ), - "python_container", - get_sample_task_metadata(), - task_interface, - custom={}, - container=get_sample_container(), - ) - sdk_promoted_task = _task.SdkTask.promote_from_model(task_template) - mock_task_fetch.return_value = sdk_promoted_task - workflow_template = get_workflow_template() - promoted_wf = _workflow_common.SdkWorkflow.promote_from_model(workflow_template) - - assert promoted_wf.interface.inputs["wf_input"] == TestPromoteExampleWf.interface.inputs["wf_input"] - assert promoted_wf.interface.outputs["wf_output_b"] == TestPromoteExampleWf.interface.outputs["wf_output_b"] - assert promoted_wf.interface.outputs["wf_output_c"] == TestPromoteExampleWf.interface.outputs["wf_output_c"] - - assert len(promoted_wf.nodes) == 1 - assert len(TestPromoteExampleWf.nodes) == 1 - assert promoted_wf.nodes[0].inputs[0] == TestPromoteExampleWf.nodes[0].inputs[0] - - def get_compiled_workflow_closure(): """ :rtype: flytekit.models.core.compiler.CompiledWorkflowClosure @@ -174,38 +108,3 @@ def get_compiled_workflow_closure(): cwc_pb.ParseFromString(fh.read()) return _compiler_model.CompiledWorkflowClosure.from_flyte_idl(cwc_pb) - - -def test_subworkflow_promote(): - cwc = get_compiled_workflow_closure() - primary = cwc.primary - sub_workflow_map = {sw.template.id: sw.template for sw in cwc.sub_workflows} - task_map = {t.template.id: t.template for t in cwc.tasks} - promoted_wf = _workflow_common.SdkWorkflow.promote_from_model(primary.template, sub_workflow_map, task_map) - - # This file that the promoted_wf reads contains the compiled workflow closure protobuf retrieved from Admin - # after registering a workflow that basically looks like the one below. - - @inputs(num=Types.Integer) - @outputs(out=Types.Integer) - @python_task - def inner_task(wf_params, num, out): - wf_params.logging.info("Running inner task... setting output to input") - out.set(num) - - @workflow_class() - class IdentityWorkflow(object): - a = Input(Types.Integer, default=5, help="Input for inner workflow") - odd_nums_task = inner_task(num=a) - task_output = Output(odd_nums_task.outputs.out, sdk_type=Types.Integer) - - @workflow_class() - class StaticSubWorkflowCaller(object): - outer_a = Input(Types.Integer, default=5, help="Input for inner workflow") - identity_wf_execution = IdentityWorkflow(a=outer_a) - wf_output = Output(identity_wf_execution.outputs.task_output, sdk_type=Types.Integer) - - assert StaticSubWorkflowCaller.interface == promoted_wf.interface - assert StaticSubWorkflowCaller.nodes[0].id == promoted_wf.nodes[0].id - assert StaticSubWorkflowCaller.nodes[0].inputs == promoted_wf.nodes[0].inputs - assert StaticSubWorkflowCaller.outputs == promoted_wf.outputs diff --git a/tests/flytekit/unit/common_tests/types/impl/__init__.py b/tests/flytekit/unit/common_tests/types/impl/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/tests/flytekit/unit/common_tests/types/impl/test_blobs.py b/tests/flytekit/unit/common_tests/types/impl/test_blobs.py deleted file mode 100644 index 3b61837e44..0000000000 --- a/tests/flytekit/unit/common_tests/types/impl/test_blobs.py +++ /dev/null @@ -1,332 +0,0 @@ -import os - -import pytest - -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.types.impl import blobs -from flytekit.common.utils import AutoDeletingTempDir -from flytekit.models.core import types as _core_types -from flytekit.sdk import test_utils - - -def test_blob(): - b = blobs.Blob("/tmp/fake") - assert b.remote_location == "/tmp/fake" - assert b.local_path is None - assert b.mode == "rb" - assert b.metadata.type.format == "" - assert b.metadata.type.dimensionality == _core_types.BlobType.BlobDimensionality.SINGLE - - -def test_blob_from_python_std(): - with test_utils.LocalTestFileSystem() as t: - with AutoDeletingTempDir("test") as wd: - tmp_name = wd.get_named_tempfile("from_python_std") - with open(tmp_name, "wb") as w: - w.write("hello hello".encode("utf-8")) - b = blobs.Blob.from_python_std(tmp_name) - assert b.mode == "wb" - assert b.metadata.type.format == "" - assert b.metadata.type.dimensionality == _core_types.BlobType.BlobDimensionality.SINGLE - assert b.remote_location.startswith(t.name) - assert b.local_path == tmp_name - with open(b.remote_location, "rb") as r: - assert r.read() == "hello hello".encode("utf-8") - - b = blobs.Blob("/tmp/fake") - b2 = blobs.Blob.from_python_std(b) - assert b == b2 - - with pytest.raises(_user_exceptions.FlyteTypeException): - blobs.Blob.from_python_std(3) - - -def test_blob_create_at(): - with test_utils.LocalTestFileSystem() as t: - with AutoDeletingTempDir("test") as wd: - tmp_name = wd.get_named_tempfile("tmp") - b = blobs.Blob.create_at_known_location(tmp_name) - assert b.local_path is None - assert b.remote_location == tmp_name - assert b.mode == "wb" - assert b.metadata.type.format == "" - assert b.metadata.type.dimensionality == _core_types.BlobType.BlobDimensionality.SINGLE - with b as w: - w.write("hello hello".encode("utf-8")) - - assert b.local_path.startswith(t.name) - with open(tmp_name, "rb") as r: - assert r.read() == "hello hello".encode("utf-8") - - -def test_blob_fetch_managed(): - with AutoDeletingTempDir("test") as wd: - with test_utils.LocalTestFileSystem() as t: - tmp_name = wd.get_named_tempfile("tmp") - with open(tmp_name, "wb") as w: - w.write("hello".encode("utf-8")) - - b = blobs.Blob.fetch(tmp_name) - assert b.local_path.startswith(t.name) - assert b.remote_location == tmp_name - assert b.mode == "rb" - assert b.metadata.type.format == "" - assert b.metadata.type.dimensionality == _core_types.BlobType.BlobDimensionality.SINGLE - with b as r: - assert r.read() == "hello".encode("utf-8") - - with pytest.raises(_user_exceptions.FlyteAssertion): - blobs.Blob.fetch(tmp_name, local_path=b.local_path) - - with open(tmp_name, "wb") as w: - w.write("bye".encode("utf-8")) - - b2 = blobs.Blob.fetch(tmp_name, local_path=b.local_path, overwrite=True) - with b2 as r: - assert r.read() == "bye".encode("utf-8") - - with pytest.raises(_user_exceptions.FlyteAssertion): - blobs.Blob.fetch(tmp_name) - - -def test_blob_fetch_unmanaged(): - with AutoDeletingTempDir("test") as wd: - with AutoDeletingTempDir("test2") as t: - tmp_name = wd.get_named_tempfile("source") - tmp_sink = t.get_named_tempfile("sink") - with open(tmp_name, "wb") as w: - w.write("hello".encode("utf-8")) - - b = blobs.Blob.fetch(tmp_name, local_path=tmp_sink) - assert b.local_path == tmp_sink - assert b.remote_location == tmp_name - assert b.mode == "rb" - assert b.metadata.type.format == "" - assert b.metadata.type.dimensionality == _core_types.BlobType.BlobDimensionality.SINGLE - with b as r: - assert r.read() == "hello".encode("utf-8") - - with pytest.raises(_user_exceptions.FlyteAssertion): - blobs.Blob.fetch(tmp_name, local_path=tmp_sink) - - with open(tmp_name, "wb") as w: - w.write("bye".encode("utf-8")) - - b2 = blobs.Blob.fetch(tmp_name, local_path=tmp_sink, overwrite=True) - with b2 as r: - assert r.read() == "bye".encode("utf-8") - - -def test_blob_double_enter(): - with test_utils.LocalTestFileSystem(): - with AutoDeletingTempDir("test") as wd: - b = blobs.Blob(wd.get_named_tempfile("sink"), mode="wb") - with b: - with pytest.raises(_user_exceptions.FlyteAssertion): - with b: - pass - - -def test_blob_download_managed(): - with AutoDeletingTempDir("test") as wd: - with test_utils.LocalTestFileSystem() as t: - tmp_name = wd.get_named_tempfile("tmp") - with open(tmp_name, "wb") as w: - w.write("hello".encode("utf-8")) - - b = blobs.Blob(tmp_name) - b.download() - assert b.local_path.startswith(t.name) - assert b.remote_location == tmp_name - assert b.mode == "rb" - assert b.metadata.type.format == "" - assert b.metadata.type.dimensionality == _core_types.BlobType.BlobDimensionality.SINGLE - with b as r: - assert r.read() == "hello".encode("utf-8") - - b2 = blobs.Blob(tmp_name) - with pytest.raises(_user_exceptions.FlyteAssertion): - b2.download(b.local_path) - - with open(tmp_name, "wb") as w: - w.write("bye".encode("utf-8")) - - b2 = blobs.Blob(tmp_name) - b2.download(local_path=b.local_path, overwrite=True) - with b2 as r: - assert r.read() == "bye".encode("utf-8") - - b = blobs.Blob(tmp_name) - with pytest.raises(_user_exceptions.FlyteAssertion): - b.download() - - -def test_blob_download_unmanaged(): - with AutoDeletingTempDir("test") as wd: - with AutoDeletingTempDir("test2") as t: - tmp_name = wd.get_named_tempfile("source") - tmp_sink = t.get_named_tempfile("sink") - with open(tmp_name, "wb") as w: - w.write("hello".encode("utf-8")) - - b = blobs.Blob(tmp_name) - b.download(tmp_sink) - assert b.local_path == tmp_sink - assert b.remote_location == tmp_name - assert b.mode == "rb" - assert b.metadata.type.format == "" - assert b.metadata.type.dimensionality == _core_types.BlobType.BlobDimensionality.SINGLE - with b as r: - assert r.read() == "hello".encode("utf-8") - - b = blobs.Blob(tmp_name) - with pytest.raises(_user_exceptions.FlyteAssertion): - b.download(tmp_sink) - - with open(tmp_name, "wb") as w: - w.write("bye".encode("utf-8")) - - b2 = blobs.Blob(tmp_name) - b2.download(tmp_sink, overwrite=True) - with b2 as r: - assert r.read() == "bye".encode("utf-8") - - -def test_multipart_blob(): - b = blobs.MultiPartBlob("/tmp/fake", mode="w", format="csv") - assert b.remote_location == "/tmp/fake/" - assert b.local_path is None - assert b.mode == "w" - assert b.metadata.type.format == "csv" - assert b.metadata.type.dimensionality == _core_types.BlobType.BlobDimensionality.MULTIPART - - -def _generate_multipart_blob_data(tmp_dir): - n = tmp_dir.get_named_tempfile("0") - with open(n, "wb") as w: - w.write("part0".encode("utf-8")) - n = tmp_dir.get_named_tempfile("1") - with open(n, "wb") as w: - w.write("part1".encode("utf-8")) - n = tmp_dir.get_named_tempfile("2") - with open(n, "wb") as w: - w.write("part2".encode("utf-8")) - - -def test_multipart_blob_from_python_std(): - with test_utils.LocalTestFileSystem() as t: - with AutoDeletingTempDir("test") as wd: - _generate_multipart_blob_data(wd) - b = blobs.MultiPartBlob.from_python_std(wd.name) - assert b.mode == "wb" - assert b.metadata.type.format == "" - assert b.metadata.type.dimensionality == _core_types.BlobType.BlobDimensionality.MULTIPART - assert b.remote_location.startswith(t.name) - assert b.local_path == wd.name - with open(os.path.join(b.remote_location, "0"), "rb") as r: - assert r.read() == "part0".encode("utf-8") - with open(os.path.join(b.remote_location, "1"), "rb") as r: - assert r.read() == "part1".encode("utf-8") - with open(os.path.join(b.remote_location, "2"), "rb") as r: - assert r.read() == "part2".encode("utf-8") - - b = blobs.MultiPartBlob("/tmp/fake/") - b2 = blobs.MultiPartBlob.from_python_std(b) - assert b == b2 - - with pytest.raises(_user_exceptions.FlyteTypeException): - blobs.MultiPartBlob.from_python_std(3) - - -def test_multipart_blob_create_at(): - with test_utils.LocalTestFileSystem(): - with AutoDeletingTempDir("test") as wd: - b = blobs.MultiPartBlob.create_at_known_location(wd.name) - assert b.local_path is None - assert b.remote_location == wd.name + "/" - assert b.mode == "wb" - assert b.metadata.type.format == "" - assert b.metadata.type.dimensionality == _core_types.BlobType.BlobDimensionality.MULTIPART - with b.create_part("0") as w: - w.write("part0".encode("utf-8")) - with b.create_part("1") as w: - w.write("part1".encode("utf-8")) - with b.create_part("2") as w: - w.write("part2".encode("utf-8")) - - with open(os.path.join(wd.name, "0"), "rb") as r: - assert r.read() == "part0".encode("utf-8") - with open(os.path.join(wd.name, "1"), "rb") as r: - assert r.read() == "part1".encode("utf-8") - with open(os.path.join(wd.name, "2"), "rb") as r: - assert r.read() == "part2".encode("utf-8") - - -def test_multipart_blob_fetch_managed(): - with AutoDeletingTempDir("test") as wd: - with test_utils.LocalTestFileSystem() as t: - _generate_multipart_blob_data(wd) - - b = blobs.MultiPartBlob.fetch(wd.name) - assert b.local_path.startswith(t.name) - assert b.remote_location == wd.name + "/" - assert b.mode == "rb" - assert b.metadata.type.format == "" - assert b.metadata.type.dimensionality == _core_types.BlobType.BlobDimensionality.MULTIPART - with b as r: - assert r[0].read() == "part0".encode("utf-8") - assert r[1].read() == "part1".encode("utf-8") - assert r[2].read() == "part2".encode("utf-8") - - with pytest.raises(_user_exceptions.FlyteAssertion): - blobs.MultiPartBlob.fetch(wd.name, local_path=b.local_path) - - with open(os.path.join(wd.name, "0"), "wb") as w: - w.write("bye".encode("utf-8")) - - b2 = blobs.MultiPartBlob.fetch(wd.name, local_path=b.local_path, overwrite=True) - with b2 as r: - assert r[0].read() == "bye".encode("utf-8") - assert r[1].read() == "part1".encode("utf-8") - assert r[2].read() == "part2".encode("utf-8") - - with pytest.raises(_user_exceptions.FlyteAssertion): - blobs.Blob.fetch(wd.name) - - -def test_multipart_blob_fetch_unmanaged(): - with AutoDeletingTempDir("test") as wd: - with AutoDeletingTempDir("test2") as t: - _generate_multipart_blob_data(wd) - tmp_sink = t.get_named_tempfile("sink") - - b = blobs.MultiPartBlob.fetch(wd.name, local_path=tmp_sink) - assert b.local_path == tmp_sink - assert b.remote_location == wd.name + "/" - assert b.mode == "rb" - assert b.metadata.type.format == "" - assert b.metadata.type.dimensionality == _core_types.BlobType.BlobDimensionality.MULTIPART - with b as r: - assert r[0].read() == "part0".encode("utf-8") - assert r[1].read() == "part1".encode("utf-8") - assert r[2].read() == "part2".encode("utf-8") - - with pytest.raises(_user_exceptions.FlyteAssertion): - blobs.MultiPartBlob.fetch(wd.name, local_path=tmp_sink) - - with open(os.path.join(wd.name, "0"), "wb") as w: - w.write("bye".encode("utf-8")) - - b2 = blobs.MultiPartBlob.fetch(wd.name, local_path=tmp_sink, overwrite=True) - with b2 as r: - assert r[0].read() == "bye".encode("utf-8") - assert r[1].read() == "part1".encode("utf-8") - assert r[2].read() == "part2".encode("utf-8") - - -def test_multipart_blob_no_enter_on_write(): - with test_utils.LocalTestFileSystem(): - b = blobs.MultiPartBlob.create_at_any_location() - with pytest.raises(_user_exceptions.FlyteAssertion): - with b: - pass diff --git a/tests/flytekit/unit/common_tests/types/impl/test_schema.py b/tests/flytekit/unit/common_tests/types/impl/test_schema.py deleted file mode 100644 index de5bba4a46..0000000000 --- a/tests/flytekit/unit/common_tests/types/impl/test_schema.py +++ /dev/null @@ -1,553 +0,0 @@ -import collections as _collections -import datetime as _datetime -import os as _os -import uuid as _uuid - -import pandas as _pd -import pytest as _pytest -import six.moves as _six_moves - -from flytekit.common import utils as _utils -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.types import blobs as _blobs -from flytekit.common.types import primitives as _primitives -from flytekit.common.types.impl import schema as _schema_impl -from flytekit.models import literals as _literal_models -from flytekit.models import types as _type_models -from flytekit.sdk import test_utils as _test_utils - - -def test_schema_type(): - _schema_impl.SchemaType() - _schema_impl.SchemaType([]) - _schema_impl.SchemaType( - [ - ("a", _primitives.Integer), - ("b", _primitives.String), - ("c", _primitives.Float), - ("d", _primitives.Boolean), - ("e", _primitives.Datetime), - ] - ) - - with _pytest.raises(ValueError): - _schema_impl.SchemaType({"a": _primitives.Integer}) - - with _pytest.raises(TypeError): - _schema_impl.SchemaType([("a", _blobs.Blob)]) - - with _pytest.raises(ValueError): - _schema_impl.SchemaType([("a", _primitives.Integer, 1)]) - - _schema_impl.SchemaType([("1", _primitives.Integer)]) - with _pytest.raises(TypeError): - _schema_impl.SchemaType([(1, _primitives.Integer)]) - - with _pytest.raises(TypeError): - _schema_impl.SchemaType([("1", [_primitives.Integer])]) - - -value_type_tuples = [ - ("abra", _primitives.Integer, [1, 2, 3, 4, 5]), - ("CADABRA", _primitives.Float, [1.0, 2.0, 3.0, 4.0, 5.0]), - ("HoCuS", _primitives.String, ["A", "B", "C", "D", "E"]), - ("Pocus", _primitives.Boolean, [True, False, True, False]), - ( - "locusts", - _primitives.Datetime, - [ - _datetime.datetime(day=1, month=1, year=2017, hour=1, minute=1, second=1, microsecond=1) - - _datetime.timedelta(days=i) - for i in _six_moves.range(5) - ], - ), -] - - -@_pytest.mark.parametrize("value_type_pair", value_type_tuples) -def test_simple_read_and_write_with_different_types(value_type_pair): - column_name, flyte_type, values = value_type_pair - values = [tuple([value]) for value in values] - schema_type = _schema_impl.SchemaType(columns=[(column_name, flyte_type)]) - - with _test_utils.LocalTestFileSystem() as sandbox: - with _utils.AutoDeletingTempDir("test") as t: - a = _schema_impl.Schema.create_at_known_location(t.name, mode="wb", schema_type=schema_type) - assert a.local_path is None - with a as writer: - for _ in _six_moves.range(5): - writer.write(_pd.DataFrame.from_records(values, columns=[column_name])) - assert a.local_path.startswith(sandbox.name) - assert a.local_path is None - - b = _schema_impl.Schema.create_at_known_location(t.name, mode="rb", schema_type=schema_type) - assert b.local_path is None - with b as reader: - for df in reader.iter_chunks(): - for check, actual in _six_moves.zip(values, df[column_name].tolist()): - assert check[0] == actual - assert reader.read() is None - reader.seek(0) - df = reader.read(concat=True) - for iter_count, actual in enumerate(df[column_name].tolist()): - assert values[iter_count % len(values)][0] == actual - assert b.local_path.startswith(sandbox.name) - assert b.local_path is None - - -def test_datetime_coercion_explicitly(): - """ - Sanity check that we're using a version of pyarrow that allows us to - truncate timestamps - """ - dt = _datetime.datetime(day=1, month=1, year=2017, hour=1, minute=1, second=1, microsecond=1) - values = [(dt,)] - df = _pd.DataFrame.from_records(values, columns=["testname"]) - assert df["testname"][0] == dt - - with _utils.AutoDeletingTempDir("test") as tmpdir: - tmpfile = tmpdir.get_named_tempfile("repro.parquet") - df.to_parquet(tmpfile, coerce_timestamps="ms", allow_truncated_timestamps=True) - df2 = _pd.read_parquet(tmpfile) - - dt2 = _datetime.datetime(day=1, month=1, year=2017, hour=1, minute=1, second=1) - assert df2["testname"][0] == dt2 - - -def test_datetime_coercion(): - values = [ - tuple( - [ - _datetime.datetime(day=1, month=1, year=2017, hour=1, minute=1, second=1, microsecond=1) - - _datetime.timedelta(days=x) - ] - ) - for x in _six_moves.range(5) - ] - schema_type = _schema_impl.SchemaType(columns=[("testname", _primitives.Datetime)]) - - with _test_utils.LocalTestFileSystem(): - with _utils.AutoDeletingTempDir("test") as t: - a = _schema_impl.Schema.create_at_known_location(t.name, mode="wb", schema_type=schema_type) - with a as writer: - for _ in _six_moves.range(5): - # us to ms coercion segfaults unless we explicitly allow truncation. - writer.write( - _pd.DataFrame.from_records(values, columns=["testname"]), - coerce_timestamps="ms", - allow_truncated_timestamps=True, - ) - - # TODO: Uncomment when segfault bug is resolved - # with _pytest.raises(Exception): - # writer.write( - # _pd.DataFrame.from_records(values, columns=['testname']), - # coerce_timestamps='ms') - - b = _schema_impl.Schema.create_at_known_location(t.name, mode="wb", schema_type=schema_type) - with b as writer: - for _ in _six_moves.range(5): - writer.write(_pd.DataFrame.from_records(values, columns=["testname"])) - - -@_pytest.mark.parametrize("value_type_pair", value_type_tuples) -def test_fetch(value_type_pair): - column_name, flyte_type, values = value_type_pair - values = [tuple([value]) for value in values] - schema_type = _schema_impl.SchemaType(columns=[(column_name, flyte_type)]) - - with _utils.AutoDeletingTempDir("test") as tmpdir: - for i in _six_moves.range(3): - _pd.DataFrame.from_records(values, columns=[column_name]).to_parquet( - tmpdir.get_named_tempfile(str(i).zfill(6)), coerce_timestamps="us" - ) - - with _utils.AutoDeletingTempDir("test2") as local_dir: - schema_obj = _schema_impl.Schema.fetch( - tmpdir.name, - local_path=local_dir.get_named_tempfile("schema_test"), - schema_type=schema_type, - ) - with schema_obj as reader: - for df in reader.iter_chunks(): - for check, actual in _six_moves.zip(values, df[column_name].tolist()): - assert check[0] == actual - assert reader.read() is None - reader.seek(0) - df = reader.read(concat=True) - for iter_count, actual in enumerate(df[column_name].tolist()): - assert values[iter_count % len(values)][0] == actual - - -@_pytest.mark.parametrize("value_type_pair", value_type_tuples) -def test_download(value_type_pair): - column_name, flyte_type, values = value_type_pair - values = [tuple([value]) for value in values] - schema_type = _schema_impl.SchemaType(columns=[(column_name, flyte_type)]) - - with _utils.AutoDeletingTempDir("test") as tmpdir: - for i in _six_moves.range(3): - _pd.DataFrame.from_records(values, columns=[column_name]).to_parquet( - tmpdir.get_named_tempfile(str(i).zfill(6)), coerce_timestamps="us" - ) - - with _utils.AutoDeletingTempDir("test2") as local_dir: - schema_obj = _schema_impl.Schema(tmpdir.name, schema_type=schema_type) - schema_obj.download(local_dir.get_named_tempfile(_uuid.uuid4().hex)) - with schema_obj as reader: - for df in reader.iter_chunks(): - for check, actual in _six_moves.zip(values, df[column_name].tolist()): - assert check[0] == actual - assert reader.read() is None - reader.seek(0) - df = reader.read(concat=True) - for iter_count, actual in enumerate(df[column_name].tolist()): - assert values[iter_count % len(values)][0] == actual - - with _pytest.raises(Exception): - schema_obj = _schema_impl.Schema(tmpdir.name, schema_type=schema_type) - schema_obj.download() - - with _test_utils.LocalTestFileSystem(): - schema_obj = _schema_impl.Schema(tmpdir.name, schema_type=schema_type) - schema_obj.download() - with schema_obj as reader: - for df in reader.iter_chunks(): - for check, actual in _six_moves.zip(values, df[column_name].tolist()): - assert check[0] == actual - assert reader.read() is None - reader.seek(0) - df = reader.read(concat=True) - for iter_count, actual in enumerate(df[column_name].tolist()): - assert values[iter_count % len(values)][0] == actual - - -def test_hive_queries(monkeypatch): - def return_deterministic_uuid(): - class FakeUUID4(object): - def __init__(self): - self.hex = "test_uuid" - - class Uuid(object): - def uuid4(self): - return FakeUUID4() - - return Uuid() - - monkeypatch.setattr(_schema_impl, "_uuid", return_deterministic_uuid()) - - all_types = _schema_impl.SchemaType( - [ - ("a", _primitives.Integer), - ("b", _primitives.String), - ("c", _primitives.Float), - ("d", _primitives.Boolean), - ("e", _primitives.Datetime), - ] - ) - - with _test_utils.LocalTestFileSystem(): - df, query = _schema_impl.Schema.create_from_hive_query( - "SELECT a, b, c, d, e FROM some_place WHERE i = 0", - stage_query="CREATE TEMPORARY TABLE some_place AS SELECT * FROM some_place_original", - known_location="s3://my_fixed_path/", - schema_type=all_types, - ) - - full_query = """ - CREATE TEMPORARY TABLE some_place AS SELECT * FROM some_place_original; - CREATE TEMPORARY TABLE test_uuid_tmp AS SELECT a, b, c, d, e FROM some_place WHERE i = 0; - CREATE EXTERNAL TABLE test_uuid LIKE test_uuid_tmp STORED AS PARQUET; - ALTER TABLE test_uuid SET LOCATION 's3://my_fixed_path/'; - INSERT OVERWRITE TABLE test_uuid - SELECT - a as a, - b as b, - CAST(c as double) c, - d as d, - e as e - FROM test_uuid_tmp; - DROP TABLE test_uuid; - """ - full_query = " ".join(full_query.split()) - query = " ".join(query.split()) - assert query == full_query - - # Test adding partition - full_query = """ - ALTER TABLE some_table ADD IF NOT EXISTS PARTITION ( - region = 'SEA', - ds = '2017-01-01' - ) LOCATION 's3://my_fixed_path/'; - ALTER TABLE some_table PARTITION ( - region = 'SEA', - ds = '2017-01-01' - ) SET LOCATION 's3://my_fixed_path/'; - """ - query = df.get_write_partition_to_hive_table_query( - "some_table", - partitions=_collections.OrderedDict([("region", "SEA"), ("ds", "2017-01-01")]), - ) - full_query = " ".join(full_query.split()) - query = " ".join(query.split()) - assert query == full_query - - -def test_partial_column_read(): - with _test_utils.LocalTestFileSystem(): - a = _schema_impl.Schema.create_at_any_location( - schema_type=_schema_impl.SchemaType([("a", _primitives.Integer), ("b", _primitives.Integer)]) - ) - with a as writer: - writer.write(_pd.DataFrame.from_dict({"a": [1, 2, 3, 4], "b": [5, 6, 7, 8]})) - - b = _schema_impl.Schema.fetch( - a.uri, - schema_type=_schema_impl.SchemaType([("a", _primitives.Integer), ("b", _primitives.Integer)]), - ) - with b as reader: - df = reader.read(columns=["b"]) - assert df.columns.values == ["b"] - assert df["b"].tolist() == [5, 6, 7, 8] - - -def test_casting(): - pass - - -def test_from_python_std(): - with _test_utils.LocalTestFileSystem(): - - def single_dataframe(): - df1 = _pd.DataFrame.from_dict({"a": [1, 2, 3, 4], "b": [5, 6, 7, 8]}) - s = _schema_impl.Schema.from_python_std( - t_value=df1, - schema_type=_schema_impl.SchemaType([("a", _primitives.Integer), ("b", _primitives.Integer)]), - ) - assert s is not None - n = _schema_impl.Schema.fetch( - s.uri, - schema_type=_schema_impl.SchemaType([("a", _primitives.Integer), ("b", _primitives.Integer)]), - ) - with n as reader: - df2 = reader.read() - assert df2.columns.values.all() == df1.columns.values.all() - assert df2["b"].tolist() == df1["b"].tolist() - - def list_of_dataframes(): - df1 = _pd.DataFrame.from_dict({"a": [1, 2, 3, 4], "b": [5, 6, 7, 8]}) - df2 = _pd.DataFrame.from_dict({"a": [9, 10, 11, 12], "b": [13, 14, 15, 16]}) - s = _schema_impl.Schema.from_python_std( - t_value=[df1, df2], - schema_type=_schema_impl.SchemaType([("a", _primitives.Integer), ("b", _primitives.Integer)]), - ) - assert s is not None - n = _schema_impl.Schema.fetch( - s.uri, - schema_type=_schema_impl.SchemaType([("a", _primitives.Integer), ("b", _primitives.Integer)]), - ) - with n as reader: - actual = [] - for df in reader.iter_chunks(): - assert df.columns.values.all() == df1.columns.values.all() - actual.extend(df["b"].tolist()) - b_val = df1["b"].tolist() - b_val.extend(df2["b"].tolist()) - assert actual == b_val - - def mixed_list(): - df1 = _pd.DataFrame.from_dict({"a": [1, 2, 3, 4], "b": [5, 6, 7, 8]}) - df2 = [1, 2, 3] - with _pytest.raises(_user_exceptions.FlyteTypeException): - _schema_impl.Schema.from_python_std( - t_value=[df1, df2], - schema_type=_schema_impl.SchemaType([("a", _primitives.Integer), ("b", _primitives.Integer)]), - ) - - def empty_list(): - s = _schema_impl.Schema.from_python_std( - t_value=[], - schema_type=_schema_impl.SchemaType([("a", _primitives.Integer), ("b", _primitives.Integer)]), - ) - assert s is not None - n = _schema_impl.Schema.fetch( - s.uri, - schema_type=_schema_impl.SchemaType([("a", _primitives.Integer), ("b", _primitives.Integer)]), - ) - with n as reader: - df = reader.read() - assert df is None - - single_dataframe() - mixed_list() - empty_list() - list_of_dataframes() - - -def test_promote_from_model_schema_type(): - m = _type_models.SchemaType( - [ - _type_models.SchemaType.SchemaColumn("a", _type_models.SchemaType.SchemaColumn.SchemaColumnType.BOOLEAN), - _type_models.SchemaType.SchemaColumn("b", _type_models.SchemaType.SchemaColumn.SchemaColumnType.DATETIME), - _type_models.SchemaType.SchemaColumn("c", _type_models.SchemaType.SchemaColumn.SchemaColumnType.DURATION), - _type_models.SchemaType.SchemaColumn("d", _type_models.SchemaType.SchemaColumn.SchemaColumnType.FLOAT), - _type_models.SchemaType.SchemaColumn("e", _type_models.SchemaType.SchemaColumn.SchemaColumnType.INTEGER), - _type_models.SchemaType.SchemaColumn("f", _type_models.SchemaType.SchemaColumn.SchemaColumnType.STRING), - ] - ) - s = _schema_impl.SchemaType.promote_from_model(m) - assert s.columns == m.columns - assert s.sdk_columns["a"].to_flyte_literal_type() == _primitives.Boolean.to_flyte_literal_type() - assert s.sdk_columns["b"].to_flyte_literal_type() == _primitives.Datetime.to_flyte_literal_type() - assert s.sdk_columns["c"].to_flyte_literal_type() == _primitives.Timedelta.to_flyte_literal_type() - assert s.sdk_columns["d"].to_flyte_literal_type() == _primitives.Float.to_flyte_literal_type() - assert s.sdk_columns["e"].to_flyte_literal_type() == _primitives.Integer.to_flyte_literal_type() - assert s.sdk_columns["f"].to_flyte_literal_type() == _primitives.String.to_flyte_literal_type() - assert s == m - - -def test_promote_from_model_schema(): - m = _literal_models.Schema( - "s3://some/place/", - _type_models.SchemaType( - [ - _type_models.SchemaType.SchemaColumn( - "a", _type_models.SchemaType.SchemaColumn.SchemaColumnType.BOOLEAN - ), - _type_models.SchemaType.SchemaColumn( - "b", _type_models.SchemaType.SchemaColumn.SchemaColumnType.DATETIME - ), - _type_models.SchemaType.SchemaColumn( - "c", _type_models.SchemaType.SchemaColumn.SchemaColumnType.DURATION - ), - _type_models.SchemaType.SchemaColumn("d", _type_models.SchemaType.SchemaColumn.SchemaColumnType.FLOAT), - _type_models.SchemaType.SchemaColumn( - "e", _type_models.SchemaType.SchemaColumn.SchemaColumnType.INTEGER - ), - _type_models.SchemaType.SchemaColumn("f", _type_models.SchemaType.SchemaColumn.SchemaColumnType.STRING), - ] - ), - ) - - s = _schema_impl.Schema.promote_from_model(m) - assert s.uri == "s3://some/place/" - assert s.type.sdk_columns["a"].to_flyte_literal_type() == _primitives.Boolean.to_flyte_literal_type() - assert s.type.sdk_columns["b"].to_flyte_literal_type() == _primitives.Datetime.to_flyte_literal_type() - assert s.type.sdk_columns["c"].to_flyte_literal_type() == _primitives.Timedelta.to_flyte_literal_type() - assert s.type.sdk_columns["d"].to_flyte_literal_type() == _primitives.Float.to_flyte_literal_type() - assert s.type.sdk_columns["e"].to_flyte_literal_type() == _primitives.Integer.to_flyte_literal_type() - assert s.type.sdk_columns["f"].to_flyte_literal_type() == _primitives.String.to_flyte_literal_type() - assert s == m - - -def test_create_at_known_location(): - with _test_utils.LocalTestFileSystem(): - with _utils.AutoDeletingTempDir("test") as wd: - b = _schema_impl.Schema.create_at_known_location(wd.name, schema_type=_schema_impl.SchemaType()) - assert b.local_path is None - assert b.remote_location == wd.name + "/" - assert b.mode == "wb" - - with b as w: - w.write(_pd.DataFrame.from_dict({"a": [1, 2, 3, 4], "b": [5, 6, 7, 8]})) - - df = _pd.read_parquet(_os.path.join(wd.name, "000000")) - assert list(df["a"]) == [1, 2, 3, 4] - assert list(df["b"]) == [5, 6, 7, 8] - - -def test_generic_schema_read(): - with _test_utils.LocalTestFileSystem(): - a = _schema_impl.Schema.create_at_any_location( - schema_type=_schema_impl.SchemaType([("a", _primitives.Integer), ("b", _primitives.Integer)]) - ) - with a as writer: - writer.write(_pd.DataFrame.from_dict({"a": [1, 2, 3, 4], "b": [5, 6, 7, 8]})) - - b = _schema_impl.Schema.fetch(a.remote_prefix, schema_type=_schema_impl.SchemaType([])) - with b as reader: - df = reader.read() - assert df.columns.values.tolist() == ["a", "b"] - assert df["a"].tolist() == [1, 2, 3, 4] - assert df["b"].tolist() == [5, 6, 7, 8] - - -def test_extra_schema_read(): - with _test_utils.LocalTestFileSystem(): - a = _schema_impl.Schema.create_at_any_location( - schema_type=_schema_impl.SchemaType([("a", _primitives.Integer), ("b", _primitives.Integer)]) - ) - with a as writer: - writer.write(_pd.DataFrame.from_dict({"a": [1, 2, 3, 4], "b": [5, 6, 7, 8]})) - - b = _schema_impl.Schema.fetch( - a.remote_prefix, - schema_type=_schema_impl.SchemaType([("a", _primitives.Integer)]), - ) - with b as reader: - df = reader.read(concat=True, truncate_extra_columns=False) - assert df.columns.values.tolist() == ["a", "b"] - assert df["a"].tolist() == [1, 2, 3, 4] - assert df["b"].tolist() == [5, 6, 7, 8] - - with b as reader: - df = reader.read(concat=True) - assert df.columns.values.tolist() == ["a"] - assert df["a"].tolist() == [1, 2, 3, 4] - - -def test_normal_schema_read_with_fastparquet(): - with _test_utils.LocalTestFileSystem(): - a = _schema_impl.Schema.create_at_any_location( - schema_type=_schema_impl.SchemaType([("a", _primitives.Integer), ("b", _primitives.Boolean)]) - ) - with a as writer: - writer.write(_pd.DataFrame.from_dict({"a": [1, 2, 3, 4], "b": [False, True, True, False]})) - - import os as _os - - original_engine = _os.getenv("PARQUET_ENGINE") - _os.environ["PARQUET_ENGINE"] = "fastparquet" - - b = _schema_impl.Schema.fetch(a.remote_prefix, schema_type=_schema_impl.SchemaType([])) - - with b as reader: - df = reader.read() - assert df["a"].tolist() == [1, 2, 3, 4] - assert _pd.api.types.is_bool_dtype(df.dtypes["b"]) - assert df["b"].tolist() == [False, True, True, False] - - if original_engine is None: - del _os.environ["PARQUET_ENGINE"] - else: - _os.environ["PARQUET_ENGINE"] = original_engine - - -def test_schema_read_consistency_between_two_engines(): - with _test_utils.LocalTestFileSystem(): - a = _schema_impl.Schema.create_at_any_location( - schema_type=_schema_impl.SchemaType([("a", _primitives.Integer), ("b", _primitives.Boolean)]) - ) - with a as writer: - writer.write(_pd.DataFrame.from_dict({"a": [1, 2, 3, 4], "b": [True, True, True, False]})) - - import os as _os - - original_engine = _os.getenv("PARQUET_ENGINE") - _os.environ["PARQUET_ENGINE"] = "fastparquet" - - b = _schema_impl.Schema.fetch(a.remote_prefix, schema_type=_schema_impl.SchemaType([])) - - with b as b_reader: - b_df = b_reader.read() - _os.environ["PARQUET_ENGINE"] = "pyarrow" - - c = _schema_impl.Schema.fetch(a.remote_prefix, schema_type=_schema_impl.SchemaType([])) - with c as c_reader: - c_df = c_reader.read() - assert b_df.equals(c_df) - - if original_engine is None: - del _os.environ["PARQUET_ENGINE"] - else: - _os.environ["PARQUET_ENGINE"] = original_engine diff --git a/tests/flytekit/unit/common_tests/types/test_blobs.py b/tests/flytekit/unit/common_tests/types/test_blobs.py deleted file mode 100644 index 3057b11c6a..0000000000 --- a/tests/flytekit/unit/common_tests/types/test_blobs.py +++ /dev/null @@ -1,141 +0,0 @@ -from flytekit.common.types import blobs -from flytekit.common.types.impl import blobs as blob_impl -from flytekit.models import literals as _literal_models -from flytekit.models.core import types as _core_types -from flytekit.sdk import test_utils - - -def test_blob_instantiator(): - b = blobs.BlobInstantiator.create_at_known_location("abc") - assert isinstance(b, blob_impl.Blob) - assert b.remote_location == "abc" - assert b.mode == "wb" - assert b.metadata.type.format == "" - - -def test_blob(): - with test_utils.LocalTestFileSystem() as t: - b = blobs.Blob() - assert isinstance(b, blob_impl.Blob) - assert b.remote_location.startswith(t.name) - assert b.mode == "wb" - assert b.metadata.type.format == "" - - b2 = blobs.Blob(b) - assert isinstance(b2, blobs.Blob) - assert b2.scalar.blob.uri == b.remote_location - assert b2.scalar.blob.metadata == b.metadata - - b3 = blobs.Blob.from_string("/a/b/c") - assert isinstance(b3, blobs.Blob) - assert b3.scalar.blob.uri == "/a/b/c" - assert b3.scalar.blob.metadata.type.format == "" - - -def test_blob_promote_from_model(): - m = _literal_models.Literal( - scalar=_literal_models.Scalar( - blob=_literal_models.Blob( - _literal_models.BlobMetadata( - _core_types.BlobType( - format="f", - dimensionality=_core_types.BlobType.BlobDimensionality.SINGLE, - ) - ), - "some/path", - ) - ) - ) - b = blobs.Blob.promote_from_model(m) - assert b.value.blob.uri == "some/path" - assert b.value.blob.metadata.type.format == "f" - assert b.value.blob.metadata.type.dimensionality == _core_types.BlobType.BlobDimensionality.SINGLE - - -def test_blob_to_python_std(): - impl = blob_impl.Blob("some/path", format="something") - b = blobs.Blob(impl).to_python_std() - assert b.metadata.type.format == "something" - assert b.metadata.type.dimensionality == _core_types.BlobType.BlobDimensionality.SINGLE - assert b.uri == "some/path" - - -def test_csv_instantiator(): - b = blobs.CsvInstantiator.create_at_known_location("abc") - assert isinstance(b, blob_impl.Blob) - assert b.remote_location == "abc" - assert b.mode == "w" - assert b.metadata.type.format == "csv" - - -def test_csv(): - with test_utils.LocalTestFileSystem() as t: - b = blobs.CSV() - assert isinstance(b, blob_impl.Blob) - assert b.remote_location.startswith(t.name) - assert b.mode == "w" - assert b.metadata.type.format == "csv" - - b2 = blobs.CSV(b) - assert isinstance(b2, blobs.Blob) - assert b2.scalar.blob.uri == b.remote_location - assert b2.scalar.blob.metadata == b.metadata - - b3 = blobs.CSV.from_string("/a/b/c") - assert isinstance(b3, blobs.Blob) - assert b3.scalar.blob.uri == "/a/b/c" - assert b3.scalar.blob.metadata.type.format == "csv" - - -def test_multipartblob_instantiator(): - b = blobs.MultiPartBlob.create_at_known_location("abc") - assert isinstance(b, blob_impl.MultiPartBlob) - assert b.remote_location == "abc" + "/" - assert b.mode == "wb" - assert b.metadata.type.format == "" - - -def test_multipartblob(): - with test_utils.LocalTestFileSystem() as t: - b = blobs.MultiPartBlob() - assert isinstance(b, blob_impl.MultiPartBlob) - assert b.remote_location.startswith(t.name) - assert b.mode == "wb" - assert b.metadata.type.format == "" - - b2 = blobs.MultiPartBlob(b) - assert isinstance(b2, blobs.MultiPartBlob) - assert b2.scalar.blob.uri == b.remote_location - assert b2.scalar.blob.metadata == b.metadata - - b3 = blobs.MultiPartBlob.from_string("/a/b/c") - assert isinstance(b3, blobs.MultiPartBlob) - assert b3.scalar.blob.uri == "/a/b/c/" - assert b3.scalar.blob.metadata.type.format == "" - - -def test_multipartcsv_instantiator(): - b = blobs.MultiPartCsvInstantiator.create_at_known_location("abc") - assert isinstance(b, blob_impl.MultiPartBlob) - assert b.remote_location == "abc" + "/" - assert b.mode == "w" - assert b.metadata.type.format == "csv" - - -def test_multipartcsv(): - with test_utils.LocalTestFileSystem() as t: - b = blobs.MultiPartCSV() - assert isinstance(b, blob_impl.MultiPartBlob) - assert b.remote_location.startswith(t.name) - assert b.mode == "w" - assert b.metadata.type.format == "csv" - - b2 = blobs.MultiPartCSV(b) - assert isinstance(b2, blobs.MultiPartCSV) - assert b2.scalar.blob.uri == b.remote_location - assert b2.scalar.blob.metadata == b.metadata - - b3 = blobs.MultiPartCSV.from_string("/a/b/c") - assert isinstance(b3, blobs.MultiPartCSV) - assert b3.scalar.blob.uri == "/a/b/c/" - assert b3.scalar.blob.metadata.type.format == "csv" diff --git a/tests/flytekit/unit/common_tests/types/test_containers.py b/tests/flytekit/unit/common_tests/types/test_containers.py deleted file mode 100644 index dcb8730b08..0000000000 --- a/tests/flytekit/unit/common_tests/types/test_containers.py +++ /dev/null @@ -1,173 +0,0 @@ -import pytest -from six.moves import range as _range - -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.types import containers, primitives -from flytekit.models import literals -from flytekit.models import types as literal_types - - -def test_list(): - list_type = containers.List(primitives.Integer) - assert list_type.to_flyte_literal_type().simple is None - assert list_type.to_flyte_literal_type().map_value_type is None - assert list_type.to_flyte_literal_type().schema is None - assert list_type.to_flyte_literal_type().collection_type.simple == literal_types.SimpleType.INTEGER - - list_value = list_type.from_python_std([1, 2, 3, 4]) - assert list_value.to_python_std() == [1, 2, 3, 4] - assert list_type.from_flyte_idl(list_value.to_flyte_idl()) == list_value - - assert list_value.collection.literals[0].scalar.primitive.integer == 1 - assert list_value.collection.literals[1].scalar.primitive.integer == 2 - assert list_value.collection.literals[2].scalar.primitive.integer == 3 - assert list_value.collection.literals[3].scalar.primitive.integer == 4 - - obj2 = list_type.from_string("[1, 2, 3,4]") - assert obj2 == list_value - - with pytest.raises(_user_exceptions.FlyteTypeException): - list_type.from_python_std(["a", "b", "c", "d"]) - - with pytest.raises(_user_exceptions.FlyteTypeException): - list_type.from_python_std([1, 2, 3, "abc"]) - - with pytest.raises(_user_exceptions.FlyteTypeException): - list_type.from_python_std(1) - - with pytest.raises(_user_exceptions.FlyteTypeException): - list_type.from_python_std([[1]]) - - with pytest.raises(_user_exceptions.FlyteTypeException): - list_type.from_string('["fdsa"]') - - with pytest.raises(_user_exceptions.FlyteTypeException): - list_type.from_string("[1, 2, 3, []]") - - with pytest.raises(_user_exceptions.FlyteTypeException): - list_type.from_string("'[\"not list json\"]'") - - with pytest.raises(_user_exceptions.FlyteTypeException): - list_type.from_string('["unclosed","list"') - - -def test_string_list(): - list_type = containers.List(primitives.String) - obj = list_type.from_string('["fdsa", "fff3", "fdsfhuie", "frfJliEILles", ""]') - assert len(obj.collection.literals) == 5 - assert obj.to_python_std() == ["fdsa", "fff3", "fdsfhuie", "frfJliEILles", ""] - - # Test that two classes of the same type are comparable - list_type_two = containers.List(primitives.String) - obj2 = list_type_two.from_string('["fdsa", "fff3", "fdsfhuie", "frfJliEILles", ""]') - assert obj == obj2 - - -def test_empty_parsing(): - list_type = containers.List(primitives.String) - obj = list_type.from_string("[]") - assert len(obj) == 0 - - # The String primitive type does not allow lists or maps to be converted - with pytest.raises(_user_exceptions.FlyteTypeException): - list_type.from_string('["fdjs", []]') - - with pytest.raises(_user_exceptions.FlyteTypeException): - list_type.from_string('["fdjs", {}]') - - -def test_nested_list(): - list_type = containers.List(containers.List(primitives.Integer)) - - assert list_type.to_flyte_literal_type().simple is None - assert list_type.to_flyte_literal_type().map_value_type is None - assert list_type.to_flyte_literal_type().schema is None - assert list_type.to_flyte_literal_type().collection_type.simple is None - assert list_type.to_flyte_literal_type().collection_type.map_value_type is None - assert list_type.to_flyte_literal_type().collection_type.schema is None - assert list_type.to_flyte_literal_type().collection_type.collection_type.simple == literal_types.SimpleType.INTEGER - - gt = [[1, 2, 3], [4, 5, 6], []] - list_value = list_type.from_python_std(gt) - assert list_value.to_python_std() == gt - assert list_type.from_flyte_idl(list_value.to_flyte_idl()) == list_value - - assert list_value.collection.literals[0].collection.literals[0].scalar.primitive.integer == 1 - assert list_value.collection.literals[0].collection.literals[1].scalar.primitive.integer == 2 - assert list_value.collection.literals[0].collection.literals[2].scalar.primitive.integer == 3 - - assert list_value.collection.literals[1].collection.literals[0].scalar.primitive.integer == 4 - assert list_value.collection.literals[1].collection.literals[1].scalar.primitive.integer == 5 - assert list_value.collection.literals[1].collection.literals[2].scalar.primitive.integer == 6 - - assert len(list_value.collection.literals[2].collection.literals) == 0 - - obj = list_type.from_string("[[1, 2, 3], [4, 5, 6]]") - assert len(obj) == 2 - assert len(obj.collection.literals[0]) == 3 - - -def test_reprs(): - list_type = containers.List(primitives.Integer) - obj = list_type.from_python_std(list(_range(3))) - assert obj.short_string() == "List(len=3, [Integer(0), Integer(1), Integer(2)])" - assert ( - obj.verbose_string() == "List(\n" - "\tlen=3,\n" - "\t[\n" - "\t\tInteger(0),\n" - "\t\tInteger(1),\n" - "\t\tInteger(2)\n" - "\t]\n" - ")" - ) - - nested_list_type = containers.List(containers.List(primitives.Integer)) - nested_obj = nested_list_type.from_python_std([list(_range(3)), list(_range(3))]) - - assert ( - nested_obj.short_string() - == "List>(len=2, [List(len=3, [Integer(0), Integer(1), Integer(2)]), " - "List(len=3, [Integer(0), Integer(1), Integer(2)])])" - ) - assert ( - nested_obj.verbose_string() == "List>(\n" - "\tlen=2,\n" - "\t[\n" - "\t\tList(\n" - "\t\t\tlen=3,\n" - "\t\t\t[\n" - "\t\t\t\tInteger(0),\n" - "\t\t\t\tInteger(1),\n" - "\t\t\t\tInteger(2)\n" - "\t\t\t]\n" - "\t\t),\n" - "\t\tList(\n" - "\t\t\tlen=3,\n" - "\t\t\t[\n" - "\t\t\t\tInteger(0),\n" - "\t\t\t\tInteger(1),\n" - "\t\t\t\tInteger(2)\n" - "\t\t\t]\n" - "\t\t)\n" - "\t]\n" - ")" - ) - - -def test_model_promotion(): - list_type = containers.List(primitives.Integer) - list_model = literals.Literal( - collection=literals.LiteralCollection( - literals=[ - literals.Literal(scalar=literals.Scalar(primitive=literals.Primitive(integer=0))), - literals.Literal(scalar=literals.Scalar(primitive=literals.Primitive(integer=1))), - literals.Literal(scalar=literals.Scalar(primitive=literals.Primitive(integer=2))), - ] - ) - ) - list_obj = list_type.promote_from_model(list_model) - assert len(list_obj.collection.literals) == 3 - assert isinstance(list_obj.collection.literals[0], primitives.Integer) - assert list_obj == list_type.from_python_std([0, 1, 2]) - assert list_obj == list_type([primitives.Integer(0), primitives.Integer(1), primitives.Integer(2)]) diff --git a/tests/flytekit/unit/common_tests/types/test_helpers.py b/tests/flytekit/unit/common_tests/types/test_helpers.py deleted file mode 100644 index dd8b45af23..0000000000 --- a/tests/flytekit/unit/common_tests/types/test_helpers.py +++ /dev/null @@ -1,59 +0,0 @@ -from flytekit.common.types import base_sdk_types as _base_sdk_types -from flytekit.common.types import helpers as _type_helpers -from flytekit.models import literals as _literals -from flytekit.models import types as _model_types -from flytekit.sdk import types as _sdk_types - - -def test_python_std_to_sdk_type(): - o = _type_helpers.python_std_to_sdk_type(_sdk_types.Types.Integer) - assert o.to_flyte_literal_type().simple == _model_types.SimpleType.INTEGER - - o = _type_helpers.python_std_to_sdk_type([_sdk_types.Types.Boolean]) - assert o.to_flyte_literal_type().collection_type.simple == _model_types.SimpleType.BOOLEAN - - -def test_get_sdk_type_from_literal_type(): - o = _type_helpers.get_sdk_type_from_literal_type(_model_types.LiteralType(simple=_model_types.SimpleType.FLOAT)) - assert o == _sdk_types.Types.Float - - -def test_infer_sdk_type_from_literal(): - o = _type_helpers.infer_sdk_type_from_literal( - _literals.Literal(scalar=_literals.Scalar(primitive=_literals.Primitive(string_value="abc"))) - ) - assert o == _sdk_types.Types.String - - o = _type_helpers.infer_sdk_type_from_literal( - _literals.Literal(scalar=_literals.Scalar(none_type=_literals.Void())) - ) - assert o is _base_sdk_types.Void - - -def test_get_sdk_value_from_literal(): - o = _type_helpers.get_sdk_value_from_literal(_literals.Literal(scalar=_literals.Scalar(none_type=_literals.Void()))) - assert o.to_python_std() is None - - o = _type_helpers.get_sdk_value_from_literal( - _literals.Literal(scalar=_literals.Scalar(none_type=_literals.Void())), - sdk_type=_sdk_types.Types.Integer, - ) - assert o.to_python_std() is None - - o = _type_helpers.get_sdk_value_from_literal( - _literals.Literal(scalar=_literals.Scalar(primitive=_literals.Primitive(integer=1))), - sdk_type=_sdk_types.Types.Integer, - ) - assert o.to_python_std() == 1 - - o = _type_helpers.get_sdk_value_from_literal( - _literals.Literal( - collection=_literals.LiteralCollection( - [ - _literals.Literal(scalar=_literals.Scalar(primitive=_literals.Primitive(integer=1))), - _literals.Literal(scalar=_literals.Scalar(none_type=_literals.Void())), - ] - ) - ) - ) - assert o.to_python_std() == [1, None] diff --git a/tests/flytekit/unit/common_tests/types/test_primitives.py b/tests/flytekit/unit/common_tests/types/test_primitives.py deleted file mode 100644 index b161c751dd..0000000000 --- a/tests/flytekit/unit/common_tests/types/test_primitives.py +++ /dev/null @@ -1,289 +0,0 @@ -import datetime - -import pytest -from dateutil import tz - -from flytekit.common.exceptions import user as user_exceptions -from flytekit.common.types import base_sdk_types, primitives -from flytekit.models import types as literal_types - - -def test_integer(): - # Check type specification - assert primitives.Integer.to_flyte_literal_type().simple == literal_types.SimpleType.INTEGER - - # Test value behavior - obj = primitives.Integer.from_python_std(1) - assert obj.to_python_std() == 1 - assert primitives.Integer.from_flyte_idl(obj.to_flyte_idl()) == obj - - for val in [ - 1.0, - "abc", - True, - False, - datetime.datetime.now(), - datetime.timedelta(seconds=1), - ]: - with pytest.raises(user_exceptions.FlyteTypeException): - primitives.Integer.from_python_std(val) - - obj = primitives.Integer.from_python_std(None) - assert obj.to_python_std() is None - assert primitives.Integer.from_flyte_idl(obj.to_flyte_idl()) == obj - - # Test string parsing - with pytest.raises(user_exceptions.FlyteTypeException): - primitives.Integer.from_string("books") - obj = primitives.Integer.from_string("299792458") - assert obj.to_python_std() == 299792458 - assert primitives.Integer.from_flyte_idl(obj.to_flyte_idl()) == obj - - assert obj.short_string() == "Integer(299792458)" - assert obj.verbose_string() == "Integer(299792458)" - - -def test_float(): - # Check type specification - assert primitives.Float.to_flyte_literal_type().simple == literal_types.SimpleType.FLOAT - - # Test value behavior - obj = primitives.Float.from_python_std(1.0) - assert obj.to_python_std() == 1.0 - assert primitives.Float.from_flyte_idl(obj.to_flyte_idl()) == obj - - for val in [ - 1, - "abc", - True, - False, - datetime.datetime.now(), - datetime.timedelta(seconds=1), - ]: - with pytest.raises(user_exceptions.FlyteTypeException): - primitives.Float.from_python_std(val) - - obj = primitives.Float.from_python_std(None) - assert obj.to_python_std() is None - assert primitives.Float.from_flyte_idl(obj.to_flyte_idl()) == obj - - # Test string parsing - with pytest.raises(user_exceptions.FlyteTypeException): - primitives.Float.from_string("lightning") - obj = primitives.Float.from_string("2.71828") - assert obj.to_python_std() == 2.71828 - assert primitives.Float.from_flyte_idl(obj.to_flyte_idl()) == obj - - assert obj.short_string() == "Float(2.71828)" - assert obj.verbose_string() == "Float(2.71828)" - - -def test_boolean(): - # Check type specification - assert primitives.Boolean.to_flyte_literal_type().simple == literal_types.SimpleType.BOOLEAN - - # Test value behavior - obj = primitives.Boolean.from_python_std(True) - assert obj.to_python_std() is True - assert primitives.Boolean.from_flyte_idl(obj.to_flyte_idl()) == obj - - for val in [1, 1.0, "abc", datetime.datetime.now(), datetime.timedelta(seconds=1)]: - with pytest.raises(user_exceptions.FlyteTypeException): - primitives.Boolean.from_python_std(val) - - obj = primitives.Boolean.from_python_std(None) - assert obj.to_python_std() is None - assert primitives.Boolean.from_flyte_idl(obj.to_flyte_idl()) == obj - - # Test string parsing - with pytest.raises(user_exceptions.FlyteTypeException): - primitives.Boolean.from_string("lightning") - obj = primitives.Boolean.from_string("false") - assert not obj.to_python_std() - assert primitives.Boolean.from_flyte_idl(obj.to_flyte_idl()) == obj - obj = primitives.Boolean.from_string("False") - assert not obj.to_python_std() - obj = primitives.Boolean.from_string("0") - assert not obj.to_python_std() - obj = primitives.Boolean.from_string("true") - assert obj.to_python_std() - obj = primitives.Boolean.from_string("True") - assert obj.to_python_std() - obj = primitives.Boolean.from_string("1") - assert obj.to_python_std() - assert primitives.Boolean.from_flyte_idl(obj.to_flyte_idl()) == obj - - assert obj.short_string() == "Boolean(True)" - assert obj.verbose_string() == "Boolean(True)" - - -def test_string(): - # Check type specification - assert primitives.String.to_flyte_literal_type().simple == literal_types.SimpleType.STRING - - # Test value behavior - obj = primitives.String.from_python_std("abc") - assert obj.to_python_std() == "abc" - assert primitives.String.from_flyte_idl(obj.to_flyte_idl()) == obj - - for val in [ - 1, - 1.0, - True, - False, - datetime.datetime.now(), - datetime.timedelta(seconds=1), - ]: - with pytest.raises(user_exceptions.FlyteTypeException): - primitives.String.from_python_std(val) - - obj = primitives.String.from_python_std(None) - assert obj.to_python_std() is None - assert primitives.String.from_flyte_idl(obj.to_flyte_idl()) == obj - - # Test string parsing - my_string = "this is a string" - obj = primitives.String.from_string(my_string) - assert obj.to_python_std() == my_string - assert primitives.String.from_flyte_idl(obj.to_flyte_idl()) == obj - - assert obj.short_string() == "String('this is a string')" - assert obj.verbose_string() == "String('this is a string')" - - with pytest.raises(user_exceptions.FlyteTypeException): - primitives.String.from_string([]) - - with pytest.raises(user_exceptions.FlyteTypeException): - primitives.String.from_string({}) - - -class UTC(datetime.tzinfo): - """UTC""" - - def utcoffset(self, dt): - return datetime.timedelta(0) - - def tzname(self, dt): - return "UTC" - - def dst(self, dt): - return datetime.timedelta(0) - - -def test_datetime(): - # Check type specification - assert primitives.Datetime.to_flyte_literal_type().simple == literal_types.SimpleType.DATETIME - - # Test value behavior - dt = datetime.datetime.now(tz=tz.UTC) - obj = primitives.Datetime.from_python_std(dt) - assert primitives.Datetime.from_flyte_idl(obj.to_flyte_idl()) == obj - assert obj.to_python_std() == dt - - # Timezone is required - with pytest.raises(user_exceptions.FlyteValueException): - primitives.Datetime.from_python_std(datetime.datetime.now()) - - for val in [1, 1.0, "abc", True, False, datetime.timedelta(seconds=1)]: - with pytest.raises(user_exceptions.FlyteTypeException): - primitives.Datetime.from_python_std(val) - - obj = primitives.Datetime.from_python_std(None) - assert obj.to_python_std() is None - assert primitives.Datetime.from_flyte_idl(obj.to_flyte_idl()) == obj - - # Test string parsing - with pytest.raises(user_exceptions.FlyteTypeException): - primitives.Datetime.from_string("not a real date") - obj = primitives.Datetime.from_string("2018-05-15 4:32pm UTC") - test_dt = datetime.datetime(2018, 5, 15, 16, 32, 0, 0, UTC()) - assert obj.short_string() == "Datetime(2018-05-15 16:32:00+00:00)" - assert obj.verbose_string() == "Datetime(2018-05-15 16:32:00+00:00)" - assert obj.to_python_std() == test_dt - assert primitives.Datetime.from_flyte_idl(obj.to_flyte_idl()) == obj - - -def test_timedelta(): - # Check type specification - assert primitives.Timedelta.to_flyte_literal_type().simple == literal_types.SimpleType.DURATION - - # Test value behavior - obj = primitives.Timedelta.from_python_std(datetime.timedelta(seconds=1)) - assert obj.to_python_std() == datetime.timedelta(seconds=1) - assert primitives.Timedelta.from_flyte_idl(obj.to_flyte_idl()) == obj - - for val in [1.0, "abc", True, False, datetime.datetime.now()]: - with pytest.raises(user_exceptions.FlyteTypeException): - primitives.Timedelta.from_python_std(val) - - obj = primitives.Timedelta.from_python_std(None) - assert obj.to_python_std() is None - assert primitives.Timedelta.from_flyte_idl(obj.to_flyte_idl()) == obj - - # Test string parsing - with pytest.raises(user_exceptions.FlyteTypeException): - primitives.Timedelta.from_string("not a real duration") - obj = primitives.Timedelta.from_string("15 hours, 1.1 second") - test_d = datetime.timedelta(hours=15, seconds=1, milliseconds=100) - assert obj.short_string() == "Timedelta(15:00:01.100000)" - assert obj.verbose_string() == "Timedelta(15:00:01.100000)" - assert obj.to_python_std() == test_d - assert primitives.Timedelta.from_flyte_idl(obj.to_flyte_idl()) == obj - - -def test_void(): - # Check type specification - with pytest.raises(user_exceptions.FlyteAssertion): - base_sdk_types.Void.to_flyte_literal_type() - - # Test value behavior - for val in [ - 1, - 1.0, - "abc", - True, - False, - datetime.datetime.now(), - datetime.timedelta(seconds=1), - None, - ]: - assert base_sdk_types.Void.from_python_std(val).to_python_std() is None - - obj = base_sdk_types.Void() - assert base_sdk_types.Void.from_flyte_idl(obj.to_flyte_idl()) == obj - - assert obj.short_string() == "Void()" - assert obj.verbose_string() == "Void()" - - -def test_generic(): - # Check type specification - assert primitives.Generic.to_flyte_literal_type().simple == literal_types.SimpleType.STRUCT - - # Test value behavior - d = {"a": [1, 2, 3], "b": "abc", "c": 1, "d": {"a": 1}} - obj = primitives.Generic.from_python_std(d) - assert obj.to_python_std() == d - assert primitives.Generic.from_flyte_idl(obj.to_flyte_idl()) == obj - - for val in [ - 1.0, - "abc", - True, - False, - datetime.datetime.now(), - datetime.timedelta(seconds=1), - ]: - with pytest.raises(user_exceptions.FlyteTypeException): - primitives.Generic.from_python_std(val) - - obj = primitives.Generic.from_python_std(None) - assert obj.to_python_std() is None - assert primitives.Generic.from_flyte_idl(obj.to_flyte_idl()) == obj - - # Test string parsing - with pytest.raises(user_exceptions.FlyteValueException): - primitives.Generic.from_string("1") - obj = primitives.Generic.from_string('{"a": 1.0}') - assert obj.to_python_std() == {"a": 1.0} - assert primitives.Generic.from_flyte_idl(obj.to_flyte_idl()) == obj diff --git a/tests/flytekit/unit/common_tests/types/test_proto.py b/tests/flytekit/unit/common_tests/types/test_proto.py deleted file mode 100644 index 074fb1dfd8..0000000000 --- a/tests/flytekit/unit/common_tests/types/test_proto.py +++ /dev/null @@ -1,63 +0,0 @@ -import base64 as _base64 - -import pytest as _pytest -from flyteidl.core import errors_pb2 as _errors_pb2 - -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.types import proto as _proto -from flytekit.common.types.proto import ProtobufType -from flytekit.models import types as _type_models - - -def test_wrong_type(): - with _pytest.raises(_user_exceptions.FlyteTypeException): - _proto.create_protobuf(int) - - -def test_proto_to_literal_type(): - proto_type = _proto.create_protobuf(_errors_pb2.ContainerError) - assert proto_type.to_flyte_literal_type().simple == _type_models.SimpleType.BINARY - assert len(proto_type.to_flyte_literal_type().metadata) == 1 - assert ( - proto_type.to_flyte_literal_type().metadata[_proto.Protobuf.PB_FIELD_KEY] - == "flyteidl.core.errors_pb2.ContainerError" - ) - - -def test_proto(): - proto_type = _proto.create_protobuf(_errors_pb2.ContainerError) - assert proto_type.short_class_string() == "Types.Proto(flyteidl.core.errors_pb2.ContainerError)" - run_test_proto_type(proto_type) - - -def test_generic_proto(): - proto_type = _proto.create_generic(_errors_pb2.ContainerError) - assert proto_type.short_class_string() == "Types.GenericProto(flyteidl.core.errors_pb2.ContainerError)" - run_test_proto_type(proto_type) - - -def run_test_proto_type(proto_type: ProtobufType): - pb = _errors_pb2.ContainerError(code="code", message="message") - obj = proto_type.from_python_std(pb) - obj2 = proto_type.from_flyte_idl(obj.to_flyte_idl()) - assert obj == obj2 - - obj = obj.to_python_std() - obj2 = obj2.to_python_std() - - assert obj.code == "code" - assert obj.message == "message" - - assert obj2.code == "code" - assert obj2.message == "message" - - -def test_from_string(): - proto_type = _proto.create_protobuf(_errors_pb2.ContainerError) - - pb = _errors_pb2.ContainerError(code="code", message="message") - pb_str = _base64.b64encode(pb.SerializeToString()) - - obj = proto_type.from_string(pb_str) - assert obj.to_python_std().code == "code" - assert obj.to_python_std().message == "message" diff --git a/tests/flytekit/unit/common_tests/types/test_schema.py b/tests/flytekit/unit/common_tests/types/test_schema.py deleted file mode 100644 index 02bfb8f55e..0000000000 --- a/tests/flytekit/unit/common_tests/types/test_schema.py +++ /dev/null @@ -1,69 +0,0 @@ -from flytekit.common.types import primitives, schema -from flytekit.common.types.impl import schema as schema_impl -from flytekit.sdk import test_utils - -_ALL_COLUMN_TYPES = [ - ("a", primitives.Integer), - ("b", primitives.String), - ("c", primitives.Float), - ("d", primitives.Datetime), - ("e", primitives.Timedelta), - ("f", primitives.Boolean), -] - - -def test_generic_schema_instantiator(): - instantiator = schema.schema_instantiator() - b = instantiator.create_at_known_location("abc") - assert isinstance(b, schema_impl.Schema) - assert b.remote_location == "abc/" - assert b.mode == "wb" - assert len(b.type.columns) == 0 - - -def test_typed_schema_instantiator(): - instantiator = schema.schema_instantiator(_ALL_COLUMN_TYPES) - b = instantiator.create_at_known_location("abc") - assert isinstance(b, schema_impl.Schema) - assert b.remote_location == "abc/" - assert b.mode == "wb" - assert len(b.type.columns) == len(_ALL_COLUMN_TYPES) - assert list(b.type.sdk_columns.items()) == _ALL_COLUMN_TYPES - - -def test_generic_schema(): - with test_utils.LocalTestFileSystem() as t: - instantiator = schema.schema_instantiator() - b = instantiator() - assert isinstance(b, schema_impl.Schema) - assert b.mode == "wb" - assert len(b.type.columns) == 0 - assert b.remote_location.startswith(t.name) - - -def test_typed_schema(): - with test_utils.LocalTestFileSystem() as t: - instantiator = schema.schema_instantiator(_ALL_COLUMN_TYPES) - b = instantiator() - assert isinstance(b, schema_impl.Schema) - assert b.mode == "wb" - assert len(b.type.columns) == len(_ALL_COLUMN_TYPES) - assert list(b.type.sdk_columns.items()) == _ALL_COLUMN_TYPES - assert b.remote_location.startswith(t.name) - - -# Ensures that subclassing types works inside a schema. -def test_casting(): - class MyDateTime(primitives.Datetime): - ... - - with test_utils.LocalTestFileSystem(): - test_columns_1 = [("altered", MyDateTime)] - test_columns_2 = [("altered", primitives.Datetime)] - - instantiator_1 = schema.schema_instantiator(test_columns_1) - a = instantiator_1() - - instantiator_2 = schema.schema_instantiator(test_columns_2) - - a.cast_to(instantiator_2._schema_type) diff --git a/tests/flytekit/unit/configuration/test_waterfall.py b/tests/flytekit/unit/configuration/test_waterfall.py index 2b35d641c0..a4cd749a80 100644 --- a/tests/flytekit/unit/configuration/test_waterfall.py +++ b/tests/flytekit/unit/configuration/test_waterfall.py @@ -1,7 +1,7 @@ import os as _os -from flytekit.common.utils import AutoDeletingTempDir as _AutoDeletingTempDir from flytekit.configuration import common as _common +from flytekit.core.utils import AutoDeletingTempDir as _AutoDeletingTempDir def test_lookup_waterfall_raw_env_var(): diff --git a/tests/flytekit/unit/contrib/__init__.py b/tests/flytekit/unit/contrib/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/tests/flytekit/unit/contrib/sensors/__init__.py b/tests/flytekit/unit/contrib/sensors/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/tests/flytekit/unit/contrib/sensors/test_impl.py b/tests/flytekit/unit/contrib/sensors/test_impl.py deleted file mode 100644 index 8382dae762..0000000000 --- a/tests/flytekit/unit/contrib/sensors/test_impl.py +++ /dev/null @@ -1,58 +0,0 @@ -import mock -from hmsclient import HMSClient -from hmsclient.genthrift.hive_metastore import ttypes as _ttypes - -from flytekit.contrib.sensors.impl import HiveFilteredPartitionSensor, HiveNamedPartitionSensor, HiveTableSensor - - -def test_HiveTableSensor(): - hive_table_sensor = HiveTableSensor(table_name="mocked_table", host="localhost", port=1234) - assert hive_table_sensor._schema == "default" - with mock.patch.object(HMSClient, "open"): - with mock.patch.object(HMSClient, "get_table"): - success, interval = hive_table_sensor._do_poll() - assert success - assert interval is None - - with mock.patch.object(HMSClient, "get_table", side_effect=_ttypes.NoSuchObjectException()): - success, interval = hive_table_sensor._do_poll() - assert not success - assert interval is None - - -def test_HiveNamedPartitionSensor(): - hive_named_partition_sensor = HiveNamedPartitionSensor( - table_name="mocked_table", partition_names=["ds=2019-10-10", "ds=2019-10-11"], host="localhost", port=1234 - ) - assert hive_named_partition_sensor._schema == "default" - with mock.patch.object(HMSClient, "open"): - with mock.patch.object(HMSClient, "get_partition_by_name"): - success, interval = hive_named_partition_sensor._do_poll() - assert success - assert interval is None - - with mock.patch.object( - HMSClient, - "get_partition_by_name", - side_effect=_ttypes.NoSuchObjectException(), - ): - success, interval = hive_named_partition_sensor._do_poll() - assert not success - assert interval is None - - -def test_HiveFilteredPartitionSensor(): - hive_filtered_partition_sensor = HiveFilteredPartitionSensor( - table_name="mocked_table", partition_filter="ds = '2019-10-10' AND region = 'NYC'", host="localhost", port=1234 - ) - assert hive_filtered_partition_sensor._schema == "default" - with mock.patch.object(HMSClient, "open"): - with mock.patch.object(HMSClient, "get_partitions_by_filter", return_value=["any"]): - success, interval = hive_filtered_partition_sensor._do_poll() - assert success - assert interval is None - - with mock.patch.object(HMSClient, "get_partitions_by_filter", return_value=[]): - success, interval = hive_filtered_partition_sensor._do_poll() - assert not success - assert interval is None diff --git a/tests/flytekit/unit/contrib/sensors/test_task.py b/tests/flytekit/unit/contrib/sensors/test_task.py deleted file mode 100644 index 0956b2aada..0000000000 --- a/tests/flytekit/unit/contrib/sensors/test_task.py +++ /dev/null @@ -1,22 +0,0 @@ -from flytekit.contrib.sensors.base_sensor import Sensor as _Sensor -from flytekit.contrib.sensors.task import sensor_task - - -class MyMockSensor(_Sensor): - def __init__(self, **kwargs): - super(MyMockSensor, self).__init__(**kwargs) - - def _do_poll(self): - """ - :rtype: (bool, Optional[datetime.timedelta]) - """ - return True, None - - -def test_sensor_works(): - @sensor_task - def my_test_task(wf_params): - return MyMockSensor() - - out = my_test_task.unit_test() - assert len(out) == 0 diff --git a/tests/flytekit/unit/core/test_conditions.py b/tests/flytekit/unit/core/test_conditions.py index b490a4f9a0..13cfcf3706 100644 --- a/tests/flytekit/unit/core/test_conditions.py +++ b/tests/flytekit/unit/core/test_conditions.py @@ -5,11 +5,11 @@ import pytest from flytekit import task, workflow -from flytekit.common.translator import get_serializable from flytekit.core import context_manager from flytekit.core.condition import conditional from flytekit.core.context_manager import Image, ImageConfig, SerializationSettings from flytekit.models.core.workflow import Node +from flytekit.tools.translator import get_serializable default_img = Image(name="default", fqn="test", tag="tag") serialization_settings = SerializationSettings( diff --git a/tests/flytekit/unit/core/test_flyte_directory.py b/tests/flytekit/unit/core/test_flyte_directory.py index db7e8c3d10..93a4d0e039 100644 --- a/tests/flytekit/unit/core/test_flyte_directory.py +++ b/tests/flytekit/unit/core/test_flyte_directory.py @@ -7,7 +7,6 @@ import pytest -from flytekit.common.exceptions.user import FlyteAssertion from flytekit.core import context_manager from flytekit.core.context_manager import ExecutionState, FlyteContextManager, Image, ImageConfig from flytekit.core.data_persistence import FileAccessProvider @@ -15,6 +14,7 @@ from flytekit.core.task import task from flytekit.core.type_engine import TypeEngine from flytekit.core.workflow import workflow +from flytekit.exceptions.user import FlyteAssertion from flytekit.models.core.types import BlobType from flytekit.models.literals import LiteralMap from flytekit.types.directory.types import FlyteDirectory, FlyteDirToMultipartBlobTransformer diff --git a/tests/flytekit/unit/core/test_flyte_pickle.py b/tests/flytekit/unit/core/test_flyte_pickle.py index 874eea27a5..c6134558b4 100644 --- a/tests/flytekit/unit/core/test_flyte_pickle.py +++ b/tests/flytekit/unit/core/test_flyte_pickle.py @@ -1,13 +1,13 @@ from collections import OrderedDict from typing import Dict, List -from flytekit.common.translator import get_serializable from flytekit.core import context_manager from flytekit.core.context_manager import Image, ImageConfig from flytekit.core.task import task from flytekit.models.core.types import BlobType from flytekit.models.literals import BlobMetadata from flytekit.models.types import LiteralType +from flytekit.tools.translator import get_serializable from flytekit.types.pickle.pickle import FlytePickle, FlytePickleTransformer default_img = Image(name="default", fqn="test", tag="tag") diff --git a/tests/flytekit/unit/core/test_imperative.py b/tests/flytekit/unit/core/test_imperative.py index 6b99c93368..f120398306 100644 --- a/tests/flytekit/unit/core/test_imperative.py +++ b/tests/flytekit/unit/core/test_imperative.py @@ -4,16 +4,16 @@ import pandas as pd import pytest -from flytekit.common.exceptions.user import FlyteValidationException -from flytekit.common.translator import get_serializable from flytekit.core import context_manager from flytekit.core.base_task import kwtypes from flytekit.core.context_manager import Image, ImageConfig from flytekit.core.launch_plan import LaunchPlan from flytekit.core.task import reference_task, task from flytekit.core.workflow import ImperativeWorkflow, get_promise, workflow +from flytekit.exceptions.user import FlyteValidationException from flytekit.extras.sqlite3.task import SQLite3Config, SQLite3Task from flytekit.models import literals as literal_models +from flytekit.tools.translator import get_serializable from flytekit.types.file import FlyteFile from flytekit.types.schema import FlyteSchema from flytekit.types.structured.structured_dataset import StructuredDatasetType diff --git a/tests/flytekit/unit/core/test_launch_plan.py b/tests/flytekit/unit/core/test_launch_plan.py index 72d10fb2d1..baa33cd356 100644 --- a/tests/flytekit/unit/core/test_launch_plan.py +++ b/tests/flytekit/unit/core/test_launch_plan.py @@ -4,7 +4,6 @@ import pytest from flyteidl.admin import launch_plan_pb2 as _launch_plan_idl -from flytekit.common.translator import get_serializable from flytekit.core import context_manager, launch_plan, notification from flytekit.core.context_manager import Image, ImageConfig from flytekit.core.schedule import CronSchedule @@ -13,6 +12,7 @@ from flytekit.models.common import Annotations, AuthRole, Labels, RawOutputDataConfig from flytekit.models.core import execution as _execution_model from flytekit.models.core import identifier as identifier_models +from flytekit.tools.translator import get_serializable default_img = Image(name="default", fqn="test", tag="tag") serialization_settings = context_manager.SerializationSettings( diff --git a/tests/flytekit/unit/core/test_map_task.py b/tests/flytekit/unit/core/test_map_task.py index 31cbceffe9..31253ce1dd 100644 --- a/tests/flytekit/unit/core/test_map_task.py +++ b/tests/flytekit/unit/core/test_map_task.py @@ -4,12 +4,12 @@ import pytest from flytekit import LaunchPlan, map_task -from flytekit.common.translator import get_serializable from flytekit.core import context_manager from flytekit.core.context_manager import Image, ImageConfig from flytekit.core.map_task import MapPythonTask from flytekit.core.task import TaskMetadata, task from flytekit.core.workflow import workflow +from flytekit.tools.translator import get_serializable @task diff --git a/tests/flytekit/unit/core/test_node_creation.py b/tests/flytekit/unit/core/test_node_creation.py index 4b1c59c479..811b2d46e5 100644 --- a/tests/flytekit/unit/core/test_node_creation.py +++ b/tests/flytekit/unit/core/test_node_creation.py @@ -5,16 +5,16 @@ import pytest from flytekit import Resources, map_task -from flytekit.common.exceptions.user import FlyteAssertion -from flytekit.common.translator import get_serializable from flytekit.core import context_manager from flytekit.core.context_manager import Image, ImageConfig from flytekit.core.dynamic_workflow_task import dynamic from flytekit.core.node_creation import create_node from flytekit.core.task import task from flytekit.core.workflow import workflow +from flytekit.exceptions.user import FlyteAssertion from flytekit.models import literals as _literal_models from flytekit.models.task import Resources as _resources_models +from flytekit.tools.translator import get_serializable def test_normal_task(): diff --git a/tests/flytekit/unit/core/test_references.py b/tests/flytekit/unit/core/test_references.py index f4fa715562..5151539236 100644 --- a/tests/flytekit/unit/core/test_references.py +++ b/tests/flytekit/unit/core/test_references.py @@ -3,7 +3,6 @@ import pytest -from flytekit.common.translator import get_serializable from flytekit.core import context_manager from flytekit.core.base_task import kwtypes from flytekit.core.context_manager import Image, ImageConfig @@ -16,6 +15,7 @@ from flytekit.core.testing import patch, task_mock from flytekit.core.workflow import reference_workflow, workflow from flytekit.models.core import identifier as _identifier_model +from flytekit.tools.translator import get_serializable # This is used for docs diff --git a/tests/flytekit/unit/core/test_resolver.py b/tests/flytekit/unit/core/test_resolver.py index 4fa33bbe34..ea44099587 100644 --- a/tests/flytekit/unit/core/test_resolver.py +++ b/tests/flytekit/unit/core/test_resolver.py @@ -3,7 +3,6 @@ import pytest -from flytekit.common.translator import get_serializable from flytekit.core import context_manager from flytekit.core.base_task import TaskResolverMixin from flytekit.core.class_based_resolver import ClassStorageTaskResolver @@ -11,6 +10,7 @@ from flytekit.core.python_auto_container import default_task_resolver from flytekit.core.task import task from flytekit.core.workflow import workflow +from flytekit.tools.translator import get_serializable default_img = Image(name="default", fqn="test", tag="tag") serialization_settings = context_manager.SerializationSettings( diff --git a/tests/flytekit/unit/core/test_serialization.py b/tests/flytekit/unit/core/test_serialization.py index fac26994a1..d3395e9fd5 100644 --- a/tests/flytekit/unit/core/test_serialization.py +++ b/tests/flytekit/unit/core/test_serialization.py @@ -5,7 +5,6 @@ import pytest from flytekit import ContainerTask, kwtypes -from flytekit.common.translator import get_serializable from flytekit.configuration import set_flyte_config_file from flytekit.core import context_manager from flytekit.core.condition import conditional @@ -13,6 +12,7 @@ from flytekit.core.task import task from flytekit.core.workflow import workflow from flytekit.models.types import SimpleType +from flytekit.tools.translator import get_serializable default_img = Image(name="default", fqn="test", tag="tag") serialization_settings = context_manager.SerializationSettings( diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 50ce9d99ac..04ea6a0bdd 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -34,6 +34,7 @@ convert_json_schema_to_python_class, dataclass_from_dict, ) +from flytekit.exceptions import user as user_exceptions from flytekit.models import types as model_types from flytekit.models.core.types import BlobType from flytekit.models.literals import Blob, BlobMetadata, Literal, LiteralCollection, LiteralMap, Primitive, Scalar, Void diff --git a/tests/flytekit/unit/core/test_type_hints.py b/tests/flytekit/unit/core/test_type_hints.py index bdb9e38c3d..367e71b467 100644 --- a/tests/flytekit/unit/core/test_type_hints.py +++ b/tests/flytekit/unit/core/test_type_hints.py @@ -37,6 +37,7 @@ from flytekit.models.interface import Parameter from flytekit.models.task import Resources as _resource_models from flytekit.models.types import LiteralType, SimpleType +from flytekit.tools.translator import get_serializable from flytekit.types.directory import FlyteDirectory, TensorboardLogs from flytekit.types.file import FlyteFile, PNGImageFile from flytekit.types.schema import FlyteSchema, SchemaOpenMode diff --git a/tests/flytekit/unit/common_tests/test_utils.py b/tests/flytekit/unit/core/test_utils.py similarity index 91% rename from tests/flytekit/unit/common_tests/test_utils.py rename to tests/flytekit/unit/core/test_utils.py index a24d585d1c..112a864b30 100644 --- a/tests/flytekit/unit/common_tests/test_utils.py +++ b/tests/flytekit/unit/core/test_utils.py @@ -1,6 +1,6 @@ import pytest -from flytekit.common.utils import _dnsify +from flytekit.core.utils import _dnsify @pytest.mark.parametrize( diff --git a/tests/flytekit/unit/core/test_workflows.py b/tests/flytekit/unit/core/test_workflows.py index 5054adc3b2..ef435393ef 100644 --- a/tests/flytekit/unit/core/test_workflows.py +++ b/tests/flytekit/unit/core/test_workflows.py @@ -6,13 +6,13 @@ from pandas.testing import assert_frame_equal from flytekit import StructuredDataset, kwtypes -from flytekit.common.exceptions.user import FlyteValidationException, FlyteValueException -from flytekit.common.translator import get_serializable from flytekit.core import context_manager from flytekit.core.condition import conditional from flytekit.core.context_manager import Image, ImageConfig from flytekit.core.task import task from flytekit.core.workflow import WorkflowFailurePolicy, WorkflowMetadata, WorkflowMetadataDefaults, workflow +from flytekit.exceptions.user import FlyteValidationException, FlyteValueException +from flytekit.tools.translator import get_serializable from flytekit.types.schema import FlyteSchema try: diff --git a/tests/flytekit/unit/engines/__init__.py b/tests/flytekit/unit/engines/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/tests/flytekit/unit/engines/flyte/__init__.py b/tests/flytekit/unit/engines/flyte/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/tests/flytekit/unit/engines/flyte/test_engine.py b/tests/flytekit/unit/engines/flyte/test_engine.py deleted file mode 100644 index ba7d478cd0..0000000000 --- a/tests/flytekit/unit/engines/flyte/test_engine.py +++ /dev/null @@ -1,807 +0,0 @@ -import os - -import pytest -from flyteidl.core import errors_pb2 -from mock import MagicMock, PropertyMock, patch - -from flytekit.common import constants, utils -from flytekit.common.exceptions import scopes -from flytekit.configuration import TemporaryConfiguration -from flytekit.engines.flyte import engine -from flytekit.models import common as _common_models -from flytekit.models import execution as _execution_models -from flytekit.models import launch_plan as _launch_plan_models -from flytekit.models import literals -from flytekit.models import task as _task_models -from flytekit.models.admin import common as _common -from flytekit.models.core import errors, identifier -from flytekit.sdk import test_utils - -_INPUT_MAP = literals.LiteralMap( - {"a": literals.Literal(scalar=literals.Scalar(primitive=literals.Primitive(integer=1)))} -) -_OUTPUT_MAP = literals.LiteralMap( - {"b": literals.Literal(scalar=literals.Scalar(primitive=literals.Primitive(integer=2)))} -) -_EMPTY_LITERAL_MAP = literals.LiteralMap(literals={}) - - -@pytest.fixture(scope="function", autouse=True) -def temp_config(): - with TemporaryConfiguration( - os.path.join( - os.path.dirname(os.path.realpath(__file__)), - "../../../common/configs/local.config", - ), - internal_overrides={ - "image": "myflyteimage:{}".format(os.environ.get("IMAGE_VERSION", "sha")), - "project": "myflyteproject", - "domain": "development", - }, - ): - yield - - -@pytest.fixture(scope="function", autouse=True) -def execution_data_locations(): - with test_utils.LocalTestFileSystem() as fs: - input_filename = fs.get_named_tempfile("inputs.pb") - output_filename = fs.get_named_tempfile("outputs.pb") - utils.write_proto_to_file(_INPUT_MAP.to_flyte_idl(), input_filename) - utils.write_proto_to_file(_OUTPUT_MAP.to_flyte_idl(), output_filename) - yield ( - _common_models.UrlBlob(input_filename, 100), - _common_models.UrlBlob(output_filename, 100), - ) - - -@scopes.system_entry_point -def _raise_system_exception(*args, **kwargs): - raise ValueError("errorERRORerror") - - -@scopes.user_entry_point -def _raise_user_exception(*args, **kwargs): - raise ValueError("userUSERuser") - - -@scopes.system_entry_point -def test_task_system_failure(): - m = MagicMock() - m.execute = _raise_system_exception - - with utils.AutoDeletingTempDir("test") as tmp: - engine.FlyteTask(m).execute(None, {"output_prefix": tmp.name}) - - doc = errors.ErrorDocument.from_flyte_idl( - utils.load_proto_from_file( - errors_pb2.ErrorDocument, - os.path.join(tmp.name, constants.ERROR_FILE_NAME), - ) - ) - assert doc.error.code == "SYSTEM:Unknown" - assert doc.error.kind == errors.ContainerError.Kind.RECOVERABLE - assert "errorERRORerror" in doc.error.message - - -@scopes.system_entry_point -def test_task_user_failure(): - m = MagicMock() - m.execute = _raise_user_exception - - with utils.AutoDeletingTempDir("test") as tmp: - engine.FlyteTask(m).execute(None, {"output_prefix": tmp.name}) - - doc = errors.ErrorDocument.from_flyte_idl( - utils.load_proto_from_file( - errors_pb2.ErrorDocument, - os.path.join(tmp.name, constants.ERROR_FILE_NAME), - ) - ) - assert doc.error.code == "USER:Unknown" - assert doc.error.kind == errors.ContainerError.Kind.NON_RECOVERABLE - assert "userUSERuser" in doc.error.message - - -@patch.object(engine._FlyteClientManager, "_CLIENT", new_callable=PropertyMock) -def test_execution_notification_overrides(mock_client_factory): - mock_client = MagicMock() - mock_client.create_execution = MagicMock(return_value=identifier.WorkflowExecutionIdentifier("xp", "xd", "xn")) - mock_client_factory.return_value = mock_client - - m = MagicMock() - type(m).id = PropertyMock( - return_value=identifier.Identifier(identifier.ResourceType.LAUNCH_PLAN, "project", "domain", "name", "version") - ) - - engine.FlyteLaunchPlan(m).launch("xp", "xd", "xn", literals.LiteralMap({}), notification_overrides=[]) - - mock_client.create_execution.assert_called_once_with( - "xp", - "xd", - "xn", - _execution_models.ExecutionSpec( - identifier.Identifier( - identifier.ResourceType.LAUNCH_PLAN, - "project", - "domain", - "name", - "version", - ), - _execution_models.ExecutionMetadata(_execution_models.ExecutionMetadata.ExecutionMode.MANUAL, "sdk", 0), - disable_all=True, - ), - literals.LiteralMap({}), - ) - - -@patch.object(engine._FlyteClientManager, "_CLIENT", new_callable=PropertyMock) -def test_execution_notification_soft_overrides(mock_client_factory): - mock_client = MagicMock() - mock_client.create_execution = MagicMock(return_value=identifier.WorkflowExecutionIdentifier("xp", "xd", "xn")) - mock_client_factory.return_value = mock_client - - m = MagicMock() - type(m).id = PropertyMock( - return_value=identifier.Identifier(identifier.ResourceType.LAUNCH_PLAN, "project", "domain", "name", "version") - ) - - notification = _common_models.Notification([0, 1, 2], email=_common_models.EmailNotification(["me@place.com"])) - - engine.FlyteLaunchPlan(m).launch("xp", "xd", "xn", literals.LiteralMap({}), notification_overrides=[notification]) - - mock_client.create_execution.assert_called_once_with( - "xp", - "xd", - "xn", - _execution_models.ExecutionSpec( - identifier.Identifier( - identifier.ResourceType.LAUNCH_PLAN, - "project", - "domain", - "name", - "version", - ), - _execution_models.ExecutionMetadata(_execution_models.ExecutionMetadata.ExecutionMode.MANUAL, "sdk", 0), - notifications=_execution_models.NotificationList([notification]), - ), - literals.LiteralMap({}), - ) - - -@patch.object(engine._FlyteClientManager, "_CLIENT", new_callable=PropertyMock) -def test_execution_label_overrides(mock_client_factory): - mock_client = MagicMock() - mock_client.create_execution = MagicMock(return_value=identifier.WorkflowExecutionIdentifier("xp", "xd", "xn")) - mock_client_factory.return_value = mock_client - - m = MagicMock() - type(m).id = PropertyMock( - return_value=identifier.Identifier(identifier.ResourceType.LAUNCH_PLAN, "project", "domain", "name", "version") - ) - - labels = _common_models.Labels({"my": "label"}) - engine.FlyteLaunchPlan(m).execute( - "xp", - "xd", - "xn", - literals.LiteralMap({}), - notification_overrides=[], - label_overrides=labels, - ) - - mock_client.create_execution.assert_called_once_with( - "xp", - "xd", - "xn", - _execution_models.ExecutionSpec( - identifier.Identifier( - identifier.ResourceType.LAUNCH_PLAN, - "project", - "domain", - "name", - "version", - ), - _execution_models.ExecutionMetadata(_execution_models.ExecutionMetadata.ExecutionMode.MANUAL, "sdk", 0), - disable_all=True, - labels=labels, - ), - literals.LiteralMap({}), - ) - - -@patch.object(engine._FlyteClientManager, "_CLIENT", new_callable=PropertyMock) -def test_execution_annotation_overrides(mock_client_factory): - mock_client = MagicMock() - mock_client.create_execution = MagicMock(return_value=identifier.WorkflowExecutionIdentifier("xp", "xd", "xn")) - mock_client_factory.return_value = mock_client - - m = MagicMock() - type(m).id = PropertyMock( - return_value=identifier.Identifier(identifier.ResourceType.LAUNCH_PLAN, "project", "domain", "name", "version") - ) - - annotations = _common_models.Annotations({"my": "annotation"}) - engine.FlyteLaunchPlan(m).launch( - "xp", - "xd", - "xn", - literals.LiteralMap({}), - notification_overrides=[], - annotation_overrides=annotations, - ) - - mock_client.create_execution.assert_called_once_with( - "xp", - "xd", - "xn", - _execution_models.ExecutionSpec( - identifier.Identifier( - identifier.ResourceType.LAUNCH_PLAN, - "project", - "domain", - "name", - "version", - ), - _execution_models.ExecutionMetadata(_execution_models.ExecutionMetadata.ExecutionMode.MANUAL, "sdk", 0), - disable_all=True, - annotations=annotations, - ), - literals.LiteralMap({}), - ) - - -@patch.object(engine._FlyteClientManager, "_CLIENT", new_callable=PropertyMock) -def test_fetch_launch_plan(mock_client_factory): - mock_client = MagicMock() - mock_client.get_launch_plan = MagicMock( - return_value=_launch_plan_models.LaunchPlan( - identifier.Identifier(identifier.ResourceType.LAUNCH_PLAN, "p1", "d1", "n1", "v1"), - MagicMock(), - MagicMock(), - ) - ) - mock_client_factory.return_value = mock_client - - lp = engine.FlyteEngineFactory().fetch_launch_plan( - identifier.Identifier(identifier.ResourceType.LAUNCH_PLAN, "p", "d", "n", "v") - ) - assert lp.id == identifier.Identifier(identifier.ResourceType.LAUNCH_PLAN, "p1", "d1", "n1", "v1") - - mock_client.get_launch_plan.assert_called_once_with( - identifier.Identifier(identifier.ResourceType.LAUNCH_PLAN, "p", "d", "n", "v") - ) - - -@patch.object(engine._FlyteClientManager, "_CLIENT", new_callable=PropertyMock) -def test_fetch_active_launch_plan(mock_client_factory): - mock_client = MagicMock() - mock_client.get_active_launch_plan = MagicMock( - return_value=_launch_plan_models.LaunchPlan( - identifier.Identifier(identifier.ResourceType.LAUNCH_PLAN, "p1", "d1", "n1", "v1"), - MagicMock(), - MagicMock(), - ) - ) - mock_client_factory.return_value = mock_client - - lp = engine.FlyteEngineFactory().fetch_launch_plan( - identifier.Identifier(identifier.ResourceType.LAUNCH_PLAN, "p", "d", "n", "") - ) - assert lp.id == identifier.Identifier(identifier.ResourceType.LAUNCH_PLAN, "p1", "d1", "n1", "v1") - - mock_client.get_active_launch_plan.assert_called_once_with(_common_models.NamedEntityIdentifier("p", "d", "n")) - - -@patch.object(engine._FlyteClientManager, "_CLIENT", new_callable=PropertyMock) -def test_get_full_execution_inputs(mock_client_factory): - mock_client = MagicMock() - mock_client.get_execution_data = MagicMock( - return_value=_execution_models.WorkflowExecutionGetDataResponse( - None, - None, - _INPUT_MAP, - _OUTPUT_MAP, - ) - ) - mock_client_factory.return_value = mock_client - - m = MagicMock() - type(m).id = PropertyMock( - return_value=identifier.WorkflowExecutionIdentifier( - "project", - "domain", - "name", - ) - ) - - inputs = engine.FlyteWorkflowExecution(m).get_inputs() - assert len(inputs.literals) == 1 - assert inputs.literals["a"].scalar.primitive.integer == 1 - mock_client.get_execution_data.assert_called_once_with( - identifier.WorkflowExecutionIdentifier("project", "domain", "name") - ) - - -@patch.object(engine._FlyteClientManager, "_CLIENT", new_callable=PropertyMock) -def test_get_execution_inputs(mock_client_factory, execution_data_locations): - mock_client = MagicMock() - mock_client.get_execution_data = MagicMock( - return_value=_execution_models.WorkflowExecutionGetDataResponse( - execution_data_locations[0], execution_data_locations[1], _EMPTY_LITERAL_MAP, _EMPTY_LITERAL_MAP - ) - ) - mock_client_factory.return_value = mock_client - - m = MagicMock() - type(m).id = PropertyMock( - return_value=identifier.WorkflowExecutionIdentifier( - "project", - "domain", - "name", - ) - ) - - inputs = engine.FlyteWorkflowExecution(m).get_inputs() - assert len(inputs.literals) == 1 - assert inputs.literals["a"].scalar.primitive.integer == 1 - mock_client.get_execution_data.assert_called_once_with( - identifier.WorkflowExecutionIdentifier("project", "domain", "name") - ) - - -@patch.object(engine._FlyteClientManager, "_CLIENT", new_callable=PropertyMock) -def test_get_full_execution_outputs(mock_client_factory): - mock_client = MagicMock() - mock_client.get_execution_data = MagicMock( - return_value=_execution_models.WorkflowExecutionGetDataResponse(None, None, _INPUT_MAP, _OUTPUT_MAP) - ) - mock_client_factory.return_value = mock_client - - m = MagicMock() - type(m).id = PropertyMock( - return_value=identifier.WorkflowExecutionIdentifier( - "project", - "domain", - "name", - ) - ) - - outputs = engine.FlyteWorkflowExecution(m).get_outputs() - assert len(outputs.literals) == 1 - assert outputs.literals["b"].scalar.primitive.integer == 2 - mock_client.get_execution_data.assert_called_once_with( - identifier.WorkflowExecutionIdentifier("project", "domain", "name") - ) - - -@patch.object(engine._FlyteClientManager, "_CLIENT", new_callable=PropertyMock) -def test_get_execution_outputs(mock_client_factory, execution_data_locations): - mock_client = MagicMock() - mock_client.get_execution_data = MagicMock( - return_value=_execution_models.WorkflowExecutionGetDataResponse( - execution_data_locations[0], execution_data_locations[1], _EMPTY_LITERAL_MAP, _EMPTY_LITERAL_MAP - ) - ) - mock_client_factory.return_value = mock_client - - m = MagicMock() - type(m).id = PropertyMock( - return_value=identifier.WorkflowExecutionIdentifier( - "project", - "domain", - "name", - ) - ) - - inputs = engine.FlyteWorkflowExecution(m).get_outputs() - assert len(inputs.literals) == 1 - assert inputs.literals["b"].scalar.primitive.integer == 2 - mock_client.get_execution_data.assert_called_once_with( - identifier.WorkflowExecutionIdentifier("project", "domain", "name") - ) - - -@patch.object(engine._FlyteClientManager, "_CLIENT", new_callable=PropertyMock) -def test_get_full_node_execution_inputs(mock_client_factory): - mock_client = MagicMock() - mock_client.get_node_execution_data = MagicMock( - return_value=_execution_models.NodeExecutionGetDataResponse( - None, - None, - _INPUT_MAP, - _OUTPUT_MAP, - ) - ) - mock_client_factory.return_value = mock_client - - m = MagicMock() - type(m).id = PropertyMock( - return_value=identifier.NodeExecutionIdentifier( - "node-a", - identifier.WorkflowExecutionIdentifier( - "project", - "domain", - "name", - ), - ) - ) - - inputs = engine.FlyteNodeExecution(m).get_inputs() - assert len(inputs.literals) == 1 - assert inputs.literals["a"].scalar.primitive.integer == 1 - mock_client.get_node_execution_data.assert_called_once_with( - identifier.NodeExecutionIdentifier( - "node-a", - identifier.WorkflowExecutionIdentifier( - "project", - "domain", - "name", - ), - ) - ) - - -@patch.object(engine._FlyteClientManager, "_CLIENT", new_callable=PropertyMock) -def test_get_node_execution_inputs(mock_client_factory, execution_data_locations): - mock_client = MagicMock() - mock_client.get_node_execution_data = MagicMock( - return_value=_execution_models.NodeExecutionGetDataResponse( - execution_data_locations[0], execution_data_locations[1], _EMPTY_LITERAL_MAP, _EMPTY_LITERAL_MAP - ) - ) - mock_client_factory.return_value = mock_client - - m = MagicMock() - type(m).id = PropertyMock( - return_value=identifier.NodeExecutionIdentifier( - "node-a", - identifier.WorkflowExecutionIdentifier( - "project", - "domain", - "name", - ), - ) - ) - - inputs = engine.FlyteNodeExecution(m).get_inputs() - assert len(inputs.literals) == 1 - assert inputs.literals["a"].scalar.primitive.integer == 1 - mock_client.get_node_execution_data.assert_called_once_with( - identifier.NodeExecutionIdentifier( - "node-a", - identifier.WorkflowExecutionIdentifier( - "project", - "domain", - "name", - ), - ) - ) - - -@patch.object(engine._FlyteClientManager, "_CLIENT", new_callable=PropertyMock) -def test_get_full_node_execution_outputs(mock_client_factory): - mock_client = MagicMock() - mock_client.get_node_execution_data = MagicMock( - return_value=_execution_models.NodeExecutionGetDataResponse(None, None, _INPUT_MAP, _OUTPUT_MAP) - ) - mock_client_factory.return_value = mock_client - - m = MagicMock() - type(m).id = PropertyMock( - return_value=identifier.NodeExecutionIdentifier( - "node-a", - identifier.WorkflowExecutionIdentifier( - "project", - "domain", - "name", - ), - ) - ) - - outputs = engine.FlyteNodeExecution(m).get_outputs() - assert len(outputs.literals) == 1 - assert outputs.literals["b"].scalar.primitive.integer == 2 - mock_client.get_node_execution_data.assert_called_once_with( - identifier.NodeExecutionIdentifier( - "node-a", - identifier.WorkflowExecutionIdentifier( - "project", - "domain", - "name", - ), - ) - ) - - -@patch.object(engine._FlyteClientManager, "_CLIENT", new_callable=PropertyMock) -def test_get_node_execution_outputs(mock_client_factory, execution_data_locations): - mock_client = MagicMock() - mock_client.get_node_execution_data = MagicMock( - return_value=_execution_models.NodeExecutionGetDataResponse( - execution_data_locations[0], execution_data_locations[1], _EMPTY_LITERAL_MAP, _EMPTY_LITERAL_MAP - ) - ) - mock_client_factory.return_value = mock_client - - m = MagicMock() - type(m).id = PropertyMock( - return_value=identifier.NodeExecutionIdentifier( - "node-a", - identifier.WorkflowExecutionIdentifier( - "project", - "domain", - "name", - ), - ) - ) - - inputs = engine.FlyteNodeExecution(m).get_outputs() - assert len(inputs.literals) == 1 - assert inputs.literals["b"].scalar.primitive.integer == 2 - mock_client.get_node_execution_data.assert_called_once_with( - identifier.NodeExecutionIdentifier( - "node-a", - identifier.WorkflowExecutionIdentifier( - "project", - "domain", - "name", - ), - ) - ) - - -@patch.object(engine._FlyteClientManager, "_CLIENT", new_callable=PropertyMock) -def test_get_full_task_execution_inputs(mock_client_factory): - mock_client = MagicMock() - mock_client.get_task_execution_data = MagicMock( - return_value=_execution_models.TaskExecutionGetDataResponse(None, None, _INPUT_MAP, _OUTPUT_MAP) - ) - mock_client_factory.return_value = mock_client - - m = MagicMock() - type(m).id = PropertyMock( - return_value=identifier.TaskExecutionIdentifier( - identifier.Identifier( - identifier.ResourceType.TASK, - "project", - "domain", - "task-name", - "version", - ), - identifier.NodeExecutionIdentifier( - "node-a", - identifier.WorkflowExecutionIdentifier( - "project", - "domain", - "name", - ), - ), - 0, - ) - ) - - inputs = engine.FlyteTaskExecution(m).get_inputs() - assert len(inputs.literals) == 1 - assert inputs.literals["a"].scalar.primitive.integer == 1 - mock_client.get_task_execution_data.assert_called_once_with( - identifier.TaskExecutionIdentifier( - identifier.Identifier( - identifier.ResourceType.TASK, - "project", - "domain", - "task-name", - "version", - ), - identifier.NodeExecutionIdentifier( - "node-a", - identifier.WorkflowExecutionIdentifier( - "project", - "domain", - "name", - ), - ), - 0, - ) - ) - - -@patch.object(engine._FlyteClientManager, "_CLIENT", new_callable=PropertyMock) -def test_get_task_execution_inputs(mock_client_factory, execution_data_locations): - mock_client = MagicMock() - mock_client.get_task_execution_data = MagicMock( - return_value=_execution_models.TaskExecutionGetDataResponse( - execution_data_locations[0], execution_data_locations[1], _EMPTY_LITERAL_MAP, _EMPTY_LITERAL_MAP - ) - ) - mock_client_factory.return_value = mock_client - - m = MagicMock() - type(m).id = PropertyMock( - return_value=identifier.TaskExecutionIdentifier( - identifier.Identifier( - identifier.ResourceType.TASK, - "project", - "domain", - "task-name", - "version", - ), - identifier.NodeExecutionIdentifier( - "node-a", - identifier.WorkflowExecutionIdentifier( - "project", - "domain", - "name", - ), - ), - 0, - ) - ) - - inputs = engine.FlyteTaskExecution(m).get_inputs() - assert len(inputs.literals) == 1 - assert inputs.literals["a"].scalar.primitive.integer == 1 - mock_client.get_task_execution_data.assert_called_once_with( - identifier.TaskExecutionIdentifier( - identifier.Identifier( - identifier.ResourceType.TASK, - "project", - "domain", - "task-name", - "version", - ), - identifier.NodeExecutionIdentifier( - "node-a", - identifier.WorkflowExecutionIdentifier( - "project", - "domain", - "name", - ), - ), - 0, - ) - ) - - -@patch.object(engine._FlyteClientManager, "_CLIENT", new_callable=PropertyMock) -def test_get_full_task_execution_outputs(mock_client_factory): - mock_client = MagicMock() - mock_client.get_task_execution_data = MagicMock( - return_value=_execution_models.TaskExecutionGetDataResponse(None, None, _INPUT_MAP, _OUTPUT_MAP) - ) - mock_client_factory.return_value = mock_client - - m = MagicMock() - type(m).id = PropertyMock( - return_value=identifier.TaskExecutionIdentifier( - identifier.Identifier( - identifier.ResourceType.TASK, - "project", - "domain", - "task-name", - "version", - ), - identifier.NodeExecutionIdentifier( - "node-a", - identifier.WorkflowExecutionIdentifier( - "project", - "domain", - "name", - ), - ), - 0, - ) - ) - - outputs = engine.FlyteTaskExecution(m).get_outputs() - assert len(outputs.literals) == 1 - assert outputs.literals["b"].scalar.primitive.integer == 2 - mock_client.get_task_execution_data.assert_called_once_with( - identifier.TaskExecutionIdentifier( - identifier.Identifier( - identifier.ResourceType.TASK, - "project", - "domain", - "task-name", - "version", - ), - identifier.NodeExecutionIdentifier( - "node-a", - identifier.WorkflowExecutionIdentifier( - "project", - "domain", - "name", - ), - ), - 0, - ) - ) - - -@patch.object(engine._FlyteClientManager, "_CLIENT", new_callable=PropertyMock) -def test_get_task_execution_outputs(mock_client_factory, execution_data_locations): - mock_client = MagicMock() - mock_client.get_task_execution_data = MagicMock( - return_value=_execution_models.TaskExecutionGetDataResponse( - execution_data_locations[0], execution_data_locations[1], _EMPTY_LITERAL_MAP, _EMPTY_LITERAL_MAP - ) - ) - mock_client_factory.return_value = mock_client - - m = MagicMock() - type(m).id = PropertyMock( - return_value=identifier.TaskExecutionIdentifier( - identifier.Identifier( - identifier.ResourceType.TASK, - "project", - "domain", - "task-name", - "version", - ), - identifier.NodeExecutionIdentifier( - "node-a", - identifier.WorkflowExecutionIdentifier( - "project", - "domain", - "name", - ), - ), - 0, - ) - ) - - inputs = engine.FlyteTaskExecution(m).get_outputs() - assert len(inputs.literals) == 1 - assert inputs.literals["b"].scalar.primitive.integer == 2 - mock_client.get_task_execution_data.assert_called_once_with( - identifier.TaskExecutionIdentifier( - identifier.Identifier( - identifier.ResourceType.TASK, - "project", - "domain", - "task-name", - "version", - ), - identifier.NodeExecutionIdentifier( - "node-a", - identifier.WorkflowExecutionIdentifier( - "project", - "domain", - "name", - ), - ), - 0, - ) - ) - - -@pytest.mark.parametrize( - "tasks", - [ - [ - _task_models.Task( - identifier.Identifier(identifier.ResourceType.TASK, "p1", "d1", "n1", "v1"), - MagicMock(), - ) - ], - [], - ], -) -@patch.object(engine._FlyteClientManager, "_CLIENT", new_callable=PropertyMock) -def test_fetch_latest_task(mock_client_factory, tasks): - mock_client = MagicMock() - mock_client.list_tasks_paginated = MagicMock(return_value=(tasks, 0)) - mock_client_factory.return_value = mock_client - - task = engine.FlyteEngineFactory().fetch_latest_task(_common_models.NamedEntityIdentifier("p", "d", "n")) - - if tasks: - assert task.id == tasks[0].id - else: - assert not task - - mock_client.list_tasks_paginated.assert_called_once_with( - _common_models.NamedEntityIdentifier("p", "d", "n"), - limit=1, - sort_by=_common.Sort("created_at", _common.Sort.Direction.DESCENDING), - ) diff --git a/tests/flytekit/unit/engines/test_loader.py b/tests/flytekit/unit/engines/test_loader.py deleted file mode 100644 index 5f6ee2d46c..0000000000 --- a/tests/flytekit/unit/engines/test_loader.py +++ /dev/null @@ -1,13 +0,0 @@ -import pytest - -from flytekit.engines import loader -from flytekit.engines.unit import engine as _unit_engine - - -def test_unit_load(): - assert isinstance(loader.get_engine("unit"), _unit_engine.UnitTestEngineFactory) - - -def test_bad_load(): - with pytest.raises(Exception): - loader.get_engine("badname") diff --git a/tests/flytekit/unit/engines/unit/__init__.py b/tests/flytekit/unit/engines/unit/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/flytekit/common/core/__init__.py b/tests/flytekit/unit/exceptions/__init__.py similarity index 100% rename from flytekit/common/core/__init__.py rename to tests/flytekit/unit/exceptions/__init__.py diff --git a/tests/flytekit/unit/common_tests/exceptions/test_base.py b/tests/flytekit/unit/exceptions/test_base.py similarity index 85% rename from tests/flytekit/unit/common_tests/exceptions/test_base.py rename to tests/flytekit/unit/exceptions/test_base.py index f4ede26b74..76b6465d8b 100644 --- a/tests/flytekit/unit/common_tests/exceptions/test_base.py +++ b/tests/flytekit/unit/exceptions/test_base.py @@ -1,4 +1,4 @@ -from flytekit.common.exceptions import base +from flytekit.exceptions import base def test_flyte_exception(): diff --git a/tests/flytekit/unit/common_tests/exceptions/test_scopes.py b/tests/flytekit/unit/exceptions/test_scopes.py similarity index 98% rename from tests/flytekit/unit/common_tests/exceptions/test_scopes.py rename to tests/flytekit/unit/exceptions/test_scopes.py index f14ced33f9..75ef74383e 100644 --- a/tests/flytekit/unit/common_tests/exceptions/test_scopes.py +++ b/tests/flytekit/unit/exceptions/test_scopes.py @@ -1,6 +1,6 @@ import pytest -from flytekit.common.exceptions import scopes, system, user +from flytekit.exceptions import scopes, system, user from flytekit.models.core import errors as _error_models diff --git a/tests/flytekit/unit/common_tests/exceptions/test_system.py b/tests/flytekit/unit/exceptions/test_system.py similarity index 98% rename from tests/flytekit/unit/common_tests/exceptions/test_system.py rename to tests/flytekit/unit/exceptions/test_system.py index d53ed00f6c..2543af320a 100644 --- a/tests/flytekit/unit/common_tests/exceptions/test_system.py +++ b/tests/flytekit/unit/exceptions/test_system.py @@ -1,4 +1,4 @@ -from flytekit.common.exceptions import base, system +from flytekit.exceptions import base, system def test_flyte_system_exception(): diff --git a/tests/flytekit/unit/common_tests/exceptions/test_user.py b/tests/flytekit/unit/exceptions/test_user.py similarity index 98% rename from tests/flytekit/unit/common_tests/exceptions/test_user.py rename to tests/flytekit/unit/exceptions/test_user.py index e3b3fbd319..78dc723def 100644 --- a/tests/flytekit/unit/common_tests/exceptions/test_user.py +++ b/tests/flytekit/unit/exceptions/test_user.py @@ -1,4 +1,4 @@ -from flytekit.common.exceptions import base, user +from flytekit.exceptions import base, user def test_flyte_user_exception(): diff --git a/tests/flytekit/unit/extras/sqlite3/test_sql_tracker.py b/tests/flytekit/unit/extras/sqlite3/test_sql_tracker.py index 0737cd5b92..b52978dd58 100644 --- a/tests/flytekit/unit/extras/sqlite3/test_sql_tracker.py +++ b/tests/flytekit/unit/extras/sqlite3/test_sql_tracker.py @@ -1,8 +1,8 @@ from collections import OrderedDict -from flytekit.common.translator import get_serializable from flytekit.core import context_manager from flytekit.core.context_manager import Image, ImageConfig +from flytekit.tools.translator import get_serializable from tests.flytekit.unit.extras.sqlite3.test_task import tk as not_tk diff --git a/tests/flytekit/unit/interfaces/data/gcs/test_gcs_proxy.py b/tests/flytekit/unit/interfaces/data/gcs/test_gcs_proxy.py deleted file mode 100644 index ea799ccf4e..0000000000 --- a/tests/flytekit/unit/interfaces/data/gcs/test_gcs_proxy.py +++ /dev/null @@ -1,80 +0,0 @@ -import os as _os - -import mock as _mock -import pytest as _pytest - -from flytekit.interfaces.data.gcs import gcs_proxy as _gcs_proxy - - -@_pytest.fixture -def mock_update_cmd_config_and_execute(): - p = _mock.patch("flytekit.interfaces.data.gcs.gcs_proxy._update_cmd_config_and_execute") - yield p.start() - p.stop() - - -@_pytest.fixture -def gsutil_parallelism(): - p = _mock.patch("flytekit.configuration.gcp.GSUTIL_PARALLELISM.get", return_value=True) - yield p.start() - p.stop() - - -@_pytest.fixture -def gcs_proxy(): - return _gcs_proxy.GCSProxy() - - -def test_upload_directory(mock_update_cmd_config_and_execute, gcs_proxy): - local_path, remote_path = "/foo/*", "gs://bar/0/" - gcs_proxy.upload_directory(local_path, remote_path) - mock_update_cmd_config_and_execute.assert_called_once_with(["gsutil", "cp", "-r", local_path, remote_path]) - - -def test_upload_directory_padding_wildcard_for_local_path(mock_update_cmd_config_and_execute, gcs_proxy): - local_path, remote_path = "/foo", "gs://bar/0/" - gcs_proxy.upload_directory(local_path, remote_path) - mock_update_cmd_config_and_execute.assert_called_once_with( - ["gsutil", "cp", "-r", _os.path.join(local_path, "*"), remote_path] - ) - - -def test_upload_directory_padding_slash_for_remote_path(mock_update_cmd_config_and_execute, gcs_proxy): - local_path, remote_path = "/foo/*", "gs://bar/0" - gcs_proxy.upload_directory(local_path, remote_path) - mock_update_cmd_config_and_execute.assert_called_once_with(["gsutil", "cp", "-r", local_path, remote_path + "/"]) - - -def test_maybe_with_gsutil_parallelism_disabled(gcs_proxy): - local_path, remote_path = "foo", "gs://bar/0/" - cmd = gcs_proxy._maybe_with_gsutil_parallelism("cp", local_path, remote_path) - assert cmd == ["gsutil", "cp", local_path, remote_path] - - -def test_maybe_with_gsutil_parallelism_enabled(gsutil_parallelism, gcs_proxy): - local_path, remote_path = "foo", "gs://bar/0/" - cmd = gcs_proxy._maybe_with_gsutil_parallelism("cp", "-r", local_path, remote_path) - assert cmd == ["gsutil", "-m", "cp", "-r", local_path, remote_path] - - -def test_download_with_parallelism(mock_update_cmd_config_and_execute, gsutil_parallelism, gcs_proxy): - local_path, remote_path = "/foo", "gs://bar/0/" - gcs_proxy.download(remote_path, local_path) - mock_update_cmd_config_and_execute.assert_called_once_with(["gsutil", "-m", "cp", remote_path, local_path]) - - -def test_upload_directory_with_parallelism(mock_update_cmd_config_and_execute, gsutil_parallelism, gcs_proxy): - local_path, remote_path = "/foo/*", "gs://bar/0/" - gcs_proxy.upload_directory(local_path, remote_path) - mock_update_cmd_config_and_execute.assert_called_once_with(["gsutil", "-m", "cp", "-r", local_path, remote_path]) - - -def test_raw_prefix_property(mock_update_cmd_config_and_execute, gsutil_parallelism, gcs_proxy): - gcs_with_raw_prefix = _gcs_proxy.GCSProxy("gcs://stuff") - assert gcs_with_raw_prefix.raw_output_data_prefix_override == "gcs://stuff" - - -def test_random_path(mock_update_cmd_config_and_execute, gsutil_parallelism, gcs_proxy): - gcs_with_raw_prefix = _gcs_proxy.GCSProxy("gcs://stuff") - result = gcs_with_raw_prefix.get_random_path() - assert result.startswith("gcs://stuff") diff --git a/tests/flytekit/unit/interfaces/data/s3/test_s3_proxy.py b/tests/flytekit/unit/interfaces/data/s3/test_s3_proxy.py deleted file mode 100644 index 7493cd3c81..0000000000 --- a/tests/flytekit/unit/interfaces/data/s3/test_s3_proxy.py +++ /dev/null @@ -1,36 +0,0 @@ -import mock as _mock - -from flytekit.interfaces.data.s3.s3proxy import AwsS3Proxy as _AwsS3Proxy - - -def test_property(): - aws = _AwsS3Proxy("s3://raw-output") - assert aws.raw_output_data_prefix_override == "s3://raw-output" - - -@_mock.patch("flytekit.configuration.aws.S3_SHARD_FORMATTER") -def test_random_path(mock_formatter): - mock_formatter.get.return_value = "s3://flyte/{}/" - - # Without raw output data prefix override - aws = _AwsS3Proxy() - p = str(aws.get_random_path()) - assert p.startswith("s3://flyte") - - # With override - aws = _AwsS3Proxy("s3://raw-output") - p = str(aws.get_random_path()) - assert p.startswith("s3://raw-output") - - -@_mock.patch("flytekit.interfaces.data.s3.s3proxy.AwsS3Proxy._check_binary") -@_mock.patch("flytekit.configuration.aws.BACKOFF_SECONDS") -@_mock.patch("flytekit.interfaces.data.s3.s3proxy._subprocess") -def test_retries(mock_subprocess, mock_delay, mock_check): - mock_delay.get.return_value = 0 - mock_subprocess.check_call.side_effect = Exception("test exception (404)") - mock_check.return_value = True - - proxy = _AwsS3Proxy() - assert proxy.exists("s3://test/fdsa/fdsa") is False - assert mock_subprocess.check_call.call_count == 4 diff --git a/tests/flytekit/unit/models/test_dynamic_spark.py b/tests/flytekit/unit/models/test_dynamic_spark.py deleted file mode 100644 index c875ef9f71..0000000000 --- a/tests/flytekit/unit/models/test_dynamic_spark.py +++ /dev/null @@ -1,29 +0,0 @@ -from flytekit.common import constants as _sdk_constants -from flytekit.configuration import TemporaryConfiguration -from flytekit.sdk import tasks as _tasks -from flytekit.sdk.types import Types as _Types - - -@_tasks.outputs(o=_Types.Integer) -@_tasks.inputs(i=_Types.Integer) -@_tasks.spark_task(spark_conf={"x": "y"}) -def my_spark_task(ctx, sc, i, o): - pass - - -@_tasks.inputs(num=_Types.Integer) -@_tasks.outputs(out=_Types.Integer) -@_tasks.dynamic_task -def spark_yield_task(wf_params, num, out): - wf_params.logging.info("Running inner task... yielding a launchplan") - t = my_spark_task.with_overrides(new_spark_conf={"a": "b"}) - o = t(i=num) - yield o - out.set(o.outputs.o) - - -def test_spark_yield(): - with TemporaryConfiguration(None, internal_overrides={"image": "fakeimage"}): - outputs = spark_yield_task.unit_test(num=1) - dj_spec = outputs[_sdk_constants.FUTURES_FILE_NAME] - print(dj_spec) diff --git a/tests/flytekit/unit/models/test_dynamic_wfs.py b/tests/flytekit/unit/models/test_dynamic_wfs.py deleted file mode 100644 index 302e87e705..0000000000 --- a/tests/flytekit/unit/models/test_dynamic_wfs.py +++ /dev/null @@ -1,125 +0,0 @@ -from flytekit.common import constants as _sdk_constants -from flytekit.sdk import tasks as _tasks -from flytekit.sdk import workflow as _workflow -from flytekit.sdk.types import Types as _Types - - -@_tasks.inputs(num=_Types.Integer) -@_tasks.outputs(out=_Types.Integer) -@_tasks.python_task -def inner_task(wf_params, num, out): - wf_params.logging.info("Running inner task... setting output to input") - out.set(num) - - -@_workflow.workflow_class() -class IdentityWorkflow(object): - a = _workflow.Input(_Types.Integer, default=5, help="Input for inner workflow") - odd_nums_task = inner_task(num=a) - task_output = _workflow.Output(odd_nums_task.outputs.out, sdk_type=_Types.Integer) - - -id_lp = IdentityWorkflow.create_launch_plan() - - -@_tasks.inputs(num=_Types.Integer) -@_tasks.outputs(out=_Types.Integer) -@_tasks.dynamic_task -def lp_yield_task(wf_params, num, out): - wf_params.logging.info("Running inner task... yielding a launchplan") - identity_lp_execution = id_lp(a=num) - yield identity_lp_execution - out.set(identity_lp_execution.outputs.task_output) - - -def test_dynamic_launch_plan_yielding(): - outputs = lp_yield_task.unit_test(num=10) - # TODO: Currently, Flytekit will not return early and not do anything if there are any workflow nodes detected - # in the output of a dynamic task. - dj_spec = outputs[_sdk_constants.FUTURES_FILE_NAME] - - assert dj_spec.min_successes == 1 - - launch_plan_node = dj_spec.nodes[0] - node_id = launch_plan_node.id - assert "models-test-dynamic-wfs-id-lp" in node_id - assert node_id.endswith("-0") - - # Assert that the output of the dynamic job spec is bound to the single node in the spec, the workflow node - # containing the launch plan - assert dj_spec.outputs[0].var == "out" - assert dj_spec.outputs[0].binding.promise.node_id == node_id - assert dj_spec.outputs[0].binding.promise.var == "task_output" - - -@_tasks.python_task -def empty_task(wf_params): - wf_params.logging.info("Running empty task") - - -@_workflow.workflow_class() -class EmptyWorkflow(object): - empty_task_task_execution = empty_task() - - -constant_workflow_lp = EmptyWorkflow.create_launch_plan() - - -@_tasks.outputs(out=_Types.Integer) -@_tasks.dynamic_task -def lp_yield_empty_wf(wf_params, out): - wf_params.logging.info("Running inner task... yielding a launchplan for empty workflow") - constant_lp_yielding_task_execution = constant_workflow_lp() - yield constant_lp_yielding_task_execution - out.set(42) - - -def test_dynamic_launch_plan_yielding_of_constant_workflow(): - outputs = lp_yield_empty_wf.unit_test() - # TODO: Currently, Flytekit will not return early and not do anything if there are any workflow nodes detected - # in the output of a dynamic task. - dj_spec = outputs[_sdk_constants.FUTURES_FILE_NAME] - - assert len(dj_spec.nodes) == 1 - assert len(dj_spec.outputs) == 1 - assert dj_spec.outputs[0].var == "out" - assert len(outputs.keys()) == 2 - - -@_tasks.inputs(num=_Types.Integer) -@_tasks.python_task -def log_only_task(wf_params, num): - wf_params.logging.info("{} was called".format(num)) - - -@_workflow.workflow_class() -class InputOnlyWorkflow(object): - a = _workflow.Input(_Types.Integer, default=5, help="Input for inner workflow") - log_only_task_execution = log_only_task(num=a) - - -input_only_workflow_lp = InputOnlyWorkflow.create_launch_plan() - - -@_tasks.dynamic_task -def lp_yield_input_only_wf(wf_params): - wf_params.logging.info("Running inner task... yielding a launchplan for input only workflow") - input_only_workflow_lp_execution = input_only_workflow_lp() - yield input_only_workflow_lp_execution - - -def test_dynamic_launch_plan_yielding_of_input_only_workflow(): - outputs = lp_yield_input_only_wf.unit_test() - # TODO: Currently, Flytekit will not return early and not do anything if there are any workflow nodes detected - # in the output of a dynamic task. - dj_spec = outputs[_sdk_constants.FUTURES_FILE_NAME] - - assert len(dj_spec.nodes) == 1 - assert len(dj_spec.outputs) == 0 - assert len(outputs.keys()) == 2 - - # Using the id of the launch plan node, and then appending /inputs.pb to the string, should give you in the outputs - # map the LiteralMap of the inputs of that node - input_key = "{}/inputs.pb".format(dj_spec.nodes[0].id) - lp_input_map = outputs[input_key] - assert lp_input_map.literals["a"] is not None diff --git a/tests/flytekit/unit/models/test_tasks.py b/tests/flytekit/unit/models/test_tasks.py index dec28f0f54..ab0fdec38a 100644 --- a/tests/flytekit/unit/models/test_tasks.py +++ b/tests/flytekit/unit/models/test_tasks.py @@ -4,7 +4,6 @@ import pytest from flyteidl.core.tasks_pb2 import TaskMetadata from google.protobuf import text_format -from k8s.io.api.core.v1 import generated_pb2 import flytekit.models.interface as interface_models import flytekit.models.literals as literal_models @@ -227,36 +226,6 @@ def test_container(resources): assert obj == task.Container.from_flyte_idl(obj.to_flyte_idl()) -def test_sidecar_task(): - pod_spec = generated_pb2.PodSpec() - container = generated_pb2.Container(name="containery") - pod_spec.containers.extend([container]) - obj = task.SidecarJob( - pod_spec=pod_spec, - primary_container_name="primary", - annotations={"a1": "a1"}, - labels={"b1": "b1"}, - ) - assert obj.primary_container_name == "primary" - assert len(obj.pod_spec.containers) == 1 - assert obj.pod_spec.containers[0].name == "containery" - assert obj.annotations["a1"] == "a1" - assert obj.labels["b1"] == "b1" - - obj2 = task.SidecarJob.from_flyte_idl(obj.to_flyte_idl()) - assert obj2 == obj - - -def test_sidecar_task_label_annotation_not_provided(): - pod_spec = generated_pb2.PodSpec() - obj = task.SidecarJob(pod_spec=pod_spec, primary_container_name="primary") - - assert obj.primary_container_name == "primary" - - obj2 = task.SidecarJob.from_flyte_idl(obj.to_flyte_idl()) - assert obj2 == obj - - def test_dataloadingconfig(): dlc = task.DataLoadingConfig( "s3://input/path", diff --git a/tests/flytekit/unit/remote/test_remote.py b/tests/flytekit/unit/remote/test_remote.py index 5904d2efdf..cd80d166e2 100644 --- a/tests/flytekit/unit/remote/test_remote.py +++ b/tests/flytekit/unit/remote/test_remote.py @@ -3,8 +3,8 @@ import pytest from mock import MagicMock, patch -from flytekit.common.exceptions import user as user_exceptions from flytekit.configuration import internal +from flytekit.exceptions import user as user_exceptions from flytekit.models import common as common_models from flytekit.models.core.identifier import ResourceType, WorkflowExecutionIdentifier from flytekit.models.execution import Execution diff --git a/tests/flytekit/unit/remote/test_wrapper_classes.py b/tests/flytekit/unit/remote/test_wrapper_classes.py index b26253e1db..f229489b14 100644 --- a/tests/flytekit/unit/remote/test_wrapper_classes.py +++ b/tests/flytekit/unit/remote/test_wrapper_classes.py @@ -3,7 +3,6 @@ import pytest -from flytekit.common.translator import gather_dependent_entities, get_serializable from flytekit.core import context_manager from flytekit.core.condition import conditional from flytekit.core.context_manager import Image, ImageConfig @@ -11,6 +10,7 @@ from flytekit.core.task import task from flytekit.core.workflow import workflow from flytekit.remote import FlyteWorkflow +from flytekit.tools.translator import gather_dependent_entities, get_serializable default_img = Image(name="default", fqn="test", tag="tag") serialization_settings = context_manager.SerializationSettings( diff --git a/tests/flytekit/unit/sdk/__init__.py b/tests/flytekit/unit/sdk/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/tests/flytekit/unit/sdk/conftest.py b/tests/flytekit/unit/sdk/conftest.py deleted file mode 100644 index 874a288263..0000000000 --- a/tests/flytekit/unit/sdk/conftest.py +++ /dev/null @@ -1,9 +0,0 @@ -import pytest as _pytest - -from flytekit.configuration import TemporaryConfiguration - - -@_pytest.fixture(scope="function", autouse=True) -def set_fake_config(): - with TemporaryConfiguration(None, internal_overrides={"image": "fakeimage"}): - yield diff --git a/tests/flytekit/unit/sdk/tasks/__init__.py b/tests/flytekit/unit/sdk/tasks/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/tests/flytekit/unit/sdk/tasks/test_dynamic_sidecar_tasks.py b/tests/flytekit/unit/sdk/tasks/test_dynamic_sidecar_tasks.py deleted file mode 100644 index 2de75de96c..0000000000 --- a/tests/flytekit/unit/sdk/tasks/test_dynamic_sidecar_tasks.py +++ /dev/null @@ -1,82 +0,0 @@ -import mock -from k8s.io.api.core.v1 import generated_pb2 - -from flytekit.common.tasks import sdk_dynamic as _sdk_dynamic -from flytekit.common.tasks import sdk_runnable as _sdk_runnable -from flytekit.common.tasks import sidecar_task as _sidecar_task -from flytekit.configuration.internal import IMAGE as _IMAGE -from flytekit.sdk.tasks import dynamic_sidecar_task, inputs, outputs, python_task -from flytekit.sdk.types import Types - - -def get_pod_spec(): - a_container = generated_pb2.Container(name="main") - a_container.command.extend(["foo", "bar"]) - a_container.volumeMounts.extend( - [ - generated_pb2.VolumeMount( - name="scratch", - mountPath="/scratch", - ) - ] - ) - - pod_spec = generated_pb2.PodSpec( - restartPolicy="Never", - ) - pod_spec.containers.extend([a_container, generated_pb2.Container(name="sidecar")]) - return pod_spec - - -with mock.patch.object(_IMAGE, "get", return_value="docker.io/blah:abc123"): - - @outputs(out1=Types.String) - @python_task - def simple_python_task(wf_params, out1): - out1.set("test") - - @inputs(in1=Types.Integer) - @outputs(out1=Types.String) - @dynamic_sidecar_task( - cpu_request="10", - memory_limit="2Gi", - environment={"foo": "bar"}, - pod_spec=get_pod_spec(), - primary_container_name="main", - ) - def simple_dynamic_sidecar_task(wf_params, in1, out1): - yield simple_python_task() - - -def test_dynamic_sidecar_task(): - assert isinstance(simple_dynamic_sidecar_task, _sdk_runnable.SdkRunnableTask) - assert isinstance(simple_dynamic_sidecar_task, _sidecar_task.SdkDynamicSidecarTask) - assert isinstance(simple_dynamic_sidecar_task, _sidecar_task.SdkSidecarTask) - assert isinstance(simple_dynamic_sidecar_task, _sdk_dynamic.SdkDynamicTaskMixin) - - pod_spec = simple_dynamic_sidecar_task.custom["podSpec"] - assert pod_spec["restartPolicy"] == "Never" - assert len(pod_spec["containers"]) == 2 - primary_container = pod_spec["containers"][0] - assert primary_container["name"] == "main" - assert primary_container["args"] == [ - "pyflyte-execute", - "--task-module", - "tests.flytekit.unit.sdk.tasks.test_dynamic_sidecar_tasks", - "--task-name", - "simple_dynamic_sidecar_task", - "--inputs", - "{{.input}}", - "--output-prefix", - "{{.outputPrefix}}", - "--raw-output-data-prefix", - "{{.rawOutputDataPrefix}}", - ] - assert primary_container["volumeMounts"] == [{"mountPath": "/scratch", "name": "scratch"}] - assert {"name": "foo", "value": "bar"} in primary_container["env"] - assert primary_container["resources"] == { - "requests": {"cpu": {"string": "10"}}, - "limits": {"memory": {"string": "2Gi"}}, - } - assert pod_spec["containers"][1]["name"] == "sidecar" - assert simple_dynamic_sidecar_task.custom["primaryContainerName"] == "main" diff --git a/tests/flytekit/unit/sdk/tasks/test_dynamic_tasks.py b/tests/flytekit/unit/sdk/tasks/test_dynamic_tasks.py deleted file mode 100644 index 44ded2e3bd..0000000000 --- a/tests/flytekit/unit/sdk/tasks/test_dynamic_tasks.py +++ /dev/null @@ -1,241 +0,0 @@ -from six import moves as _six_moves - -from flytekit.common.tasks import sdk_dynamic as _sdk_dynamic -from flytekit.common.tasks import sdk_runnable as _sdk_runnable -from flytekit.sdk.tasks import dynamic_task, inputs, outputs, python_task -from flytekit.sdk.types import Types -from flytekit.sdk.workflow import Input, Output, workflow - - -@inputs(in1=Types.Integer) -@outputs(out_str=[Types.String], out_ints=[[Types.Integer]]) -@dynamic_task -def sample_batch_task(wf_params, in1, out_str, out_ints): - res = ["I'm the first result"] - for i in _six_moves.range(0, in1): - task = sub_task(in1=i) - yield task - res.append(task.outputs.out1) - res.append("I'm after each sub-task result") - res.append("I'm the last result") - - res2 = [] - for i in _six_moves.range(0, in1): - task = int_sub_task(in1=i) - yield task - res2.append(task.outputs.out1) - - # Nested batch tasks - task = sample_batch_task_sq() - yield task - res2.append(task.outputs.out_ints) - - task = sample_batch_task_sq() - yield task - res2.append(task.outputs.out_ints) - - out_str.set(res) - out_ints.set(res2) - - -@outputs(out_ints=[Types.Integer]) -@dynamic_task -def sample_batch_task_sq(wf_params, out_ints): - res2 = [] - for i in _six_moves.range(0, 3): - task = sq_sub_task(in1=i) - yield task - res2.append(task.outputs.out1) - out_ints.set(res2) - - -@outputs(out_str=[Types.String], out_ints=[[Types.Integer]]) -@dynamic_task -def sample_batch_task_no_inputs(wf_params, out_str, out_ints): - res = ["I'm the first result"] - for i in _six_moves.range(0, 3): - task = sub_task(in1=i) - yield task - res.append(task.outputs.out1) - res.append("I'm after each sub-task result") - res.append("I'm the last result") - - res2 = [] - for i in _six_moves.range(0, 3): - task = int_sub_task(in1=i) - yield task - res2.append(task.outputs.out1) - - # Nested batch tasks - task = sample_batch_task_sq() - yield task - res2.append(task.outputs.out_ints) - - task = sample_batch_task_sq() - yield task - res2.append(task.outputs.out_ints) - - out_str.set(res) - out_ints.set(res2) - - -@inputs(in1=Types.Integer) -@outputs(out1=Types.String) -@python_task -def sub_task(wf_params, in1, out1): - out1.set("hello {}".format(in1)) - - -@inputs(in1=Types.Integer) -@outputs(out1=[Types.Integer]) -@python_task -def int_sub_task(wf_params, in1, out1): - wf_params.stats.incr("int_sub_task") - out1.set([in1, in1 * 2, in1 * 3]) - - -@inputs(in1=Types.Integer) -@outputs(out1=Types.Integer) -@python_task -def sq_sub_task(wf_params, in1, out1): - out1.set(in1 * in1) - - -@inputs(in1=Types.Integer) -@outputs(out_str=[Types.String]) -@dynamic_task -def no_future_batch_task(wf_params, in1, out_str): - out_str.set(["res1", "res2"]) - - -def manual_assign_name(): - pass - - -@inputs(task_input_num=Types.Integer) -@outputs(out=Types.Integer) -@dynamic_task -def dynamic_wf_task(wf_params, task_input_num, out): - wf_params.logging.info("Running inner task... yielding a code generated sub workflow") - - input_a = Input(Types.Integer, help="Tell me something") - node1 = sq_sub_task(in1=input_a) - - MyUnregisteredWorkflow = workflow( - inputs={"a": input_a}, - outputs={"ooo": Output(node1.outputs.out1, sdk_type=Types.Integer, help="This is an integer output")}, - nodes={"node_one": node1}, - ) - - setattr(MyUnregisteredWorkflow, "auto_assign_name", manual_assign_name) - MyUnregisteredWorkflow._platform_valid_name = "unregistered" - - unregistered_workflow_execution = MyUnregisteredWorkflow(a=task_input_num) - out.set(unregistered_workflow_execution.outputs.ooo) - - -def test_batch_task(): - assert isinstance(sample_batch_task, _sdk_runnable.SdkRunnableTask) - assert isinstance(sample_batch_task, _sdk_dynamic.SdkDynamicTask) - assert isinstance(sample_batch_task, _sdk_dynamic.SdkDynamicTaskMixin) - - expected = { - "out_str": [ - "I'm the first result", - "hello 0", - "I'm after each sub-task result", - "hello 1", - "I'm after each sub-task result", - "hello 2", - "I'm after each sub-task result", - "I'm the last result", - ], - "out_ints": [[0, 0, 0], [1, 2, 3], [2, 4, 6], [0, 1, 4], [0, 1, 4]], - } - - res = sample_batch_task.unit_test(in1=3) - assert expected == res - - -def test_no_future_batch_task(): - expected = {"out_str": ["res1", "res2"]} - - res = no_future_batch_task.unit_test(in1=3) - assert expected == res - - -def test_dynamic_workflow(): - res = dynamic_wf_task.unit_test(task_input_num=2) - dynamic_spec = res["futures.pb"] - assert len(dynamic_spec.nodes) == 1 - assert len(dynamic_spec.subworkflows) == 1 - assert len(dynamic_spec.tasks) == 1 - - -@inputs(task_input_num=Types.Integer) -@outputs(out=Types.Integer) -@dynamic_task -def nested_dynamic_wf_task(wf_params, task_input_num, out): - wf_params.logging.info("Running inner task... yielding a code generated sub workflow") - - # Inner workflow - input_a = Input(Types.Integer, help="Tell me something") - node1 = sq_sub_task(in1=input_a) - - MyUnregisteredWorkflowInner = workflow( - inputs={"a": input_a}, - outputs={"ooo": Output(node1.outputs.out1, sdk_type=Types.Integer, help="This is an integer output")}, - nodes={"node_one": node1}, - ) - - setattr(MyUnregisteredWorkflowInner, "auto_assign_name", manual_assign_name) - MyUnregisteredWorkflowInner._platform_valid_name = "unregistered" - - # Output workflow - input_a = Input(Types.Integer, help="Tell me something") - node1 = MyUnregisteredWorkflowInner(a=task_input_num) - - MyUnregisteredWorkflowOuter = workflow( - inputs={"a": input_a}, - outputs={"ooo": Output(node1.outputs.ooo, sdk_type=Types.Integer, help="This is an integer output")}, - nodes={"node_one": node1}, - ) - - setattr(MyUnregisteredWorkflowOuter, "auto_assign_name", manual_assign_name) - MyUnregisteredWorkflowOuter._platform_valid_name = "unregistered" - - unregistered_workflow_execution = MyUnregisteredWorkflowOuter(a=task_input_num) - out.set(unregistered_workflow_execution.outputs.ooo) - - -def test_nested_dynamic_workflow(): - res = nested_dynamic_wf_task.unit_test(task_input_num=2) - dynamic_spec = res["futures.pb"] - assert len(dynamic_spec.nodes) == 1 - assert len(dynamic_spec.subworkflows) == 2 - assert len(dynamic_spec.tasks) == 1 - - -@inputs(task_input_num=Types.Integer) -@dynamic_task -def dynamic_wf_no_outputs_task(wf_params, task_input_num): - wf_params.logging.info("Running inner task... yielding a code generated sub workflow") - - input_a = Input(Types.Integer, help="Tell me something") - node1 = sq_sub_task(in1=input_a) - - MyUnregisteredWorkflow = workflow(inputs={"a": input_a}, outputs={}, nodes={"node_one": node1}) - - setattr(MyUnregisteredWorkflow, "auto_assign_name", manual_assign_name) - MyUnregisteredWorkflow._platform_valid_name = "unregistered" - - unregistered_workflow_execution = MyUnregisteredWorkflow(a=task_input_num) - yield unregistered_workflow_execution - - -def test_dynamic_workflow_no_outputs(): - res = dynamic_wf_no_outputs_task.unit_test(task_input_num=2) - dynamic_spec = res["futures.pb"] - assert len(dynamic_spec.nodes) == 1 - assert len(dynamic_spec.subworkflows) == 1 - assert len(dynamic_spec.tasks) == 1 diff --git a/tests/flytekit/unit/sdk/tasks/test_hive_tasks.py b/tests/flytekit/unit/sdk/tasks/test_hive_tasks.py deleted file mode 100644 index e81edec5de..0000000000 --- a/tests/flytekit/unit/sdk/tasks/test_hive_tasks.py +++ /dev/null @@ -1,141 +0,0 @@ -import logging as _logging -from datetime import datetime as _datetime - -import six as _six - -from flytekit.common import utils as _common_utils -from flytekit.common.tasks import hive_task as _hive_task -from flytekit.common.tasks import output as _task_output -from flytekit.common.tasks import sdk_runnable as _sdk_runnable -from flytekit.common.types import base_sdk_types as _base_sdk_types -from flytekit.common.types import containers as _containers -from flytekit.common.types import helpers as _type_helpers -from flytekit.common.types import schema as _schema -from flytekit.common.types.impl.schema import Schema -from flytekit.engines import common as _common_engine -from flytekit.models import literals as _literals -from flytekit.models.core.identifier import WorkflowExecutionIdentifier -from flytekit.sdk.tasks import hive_task, inputs, outputs, qubole_hive_task -from flytekit.sdk.types import Types - - -@hive_task(cache_version="1") -def sample_hive_task_no_input(wf_params): - return _six.text_type("select 5") - - -@inputs(in1=Types.Integer) -@hive_task(cache_version="1") -def sample_hive_task(wf_params, in1): - return _six.text_type("select ") + _six.text_type(in1) - - -@hive_task -def sample_hive_task_no_queries(wf_params): - return [] - - -@qubole_hive_task( - cache_version="1", - cluster_label=_six.text_type("cluster_label"), - tags=[], -) -def sample_qubole_hive_task_no_input(wf_params): - return _six.text_type("select 5") - - -@inputs(in1=Types.Integer) -@qubole_hive_task( - cache_version="1", - cluster_label=_six.text_type("cluster_label"), - tags=[_six.text_type("tag1")], -) -def sample_qubole_hive_task(wf_params, in1): - return _six.text_type("select ") + _six.text_type(in1) - - -def test_hive_task(): - assert isinstance(sample_hive_task, _sdk_runnable.SdkRunnableTask) - assert isinstance(sample_hive_task, _hive_task.SdkHiveTask) - - sample_hive_task.unit_test(in1=5) - - -@outputs(hive_results=[Types.Schema()]) -@qubole_hive_task -def two_queries(wf_params, hive_results): - q1 = "SELECT 1" - q2 = "SELECT 'two'" - schema_1, formatted_query_1 = Schema.create_from_hive_query(select_query=q1) - schema_2, formatted_query_2 = Schema.create_from_hive_query(select_query=q2) - - hive_results.set([schema_1, schema_2]) - return [formatted_query_1, formatted_query_2] - - -def test_interface_setup(): - outs = two_queries.interface.outputs - assert outs["hive_results"].type.collection_type is not None - assert outs["hive_results"].type.collection_type.schema is not None - assert outs["hive_results"].type.collection_type.schema.columns == [] - - -def test_sdk_output_references_construction(): - references = { - name: _task_output.OutputReference(_type_helpers.get_sdk_type_from_literal_type(variable.type)) - for name, variable in _six.iteritems(two_queries.interface.outputs) - } - # Before user code is run, the outputs passed to the user code should not have values - assert references["hive_results"].sdk_value == _base_sdk_types.Void() - - # Should be a list of schemas - assert isinstance(references["hive_results"].sdk_type, _containers.TypedCollectionType) - assert isinstance(references["hive_results"].sdk_type.sub_type, _schema.SchemaInstantiator) - - -def test_hive_task_query_generation(): - with _common_utils.AutoDeletingTempDir("user_dir") as user_working_directory: - context = _common_engine.EngineContext( - execution_id=WorkflowExecutionIdentifier(project="unit_test", domain="unit_test", name="unit_test"), - execution_date=_datetime.utcnow(), - stats=None, # TODO: A mock stats object that we can read later. - logging=_logging, # TODO: A mock logging object that we can read later. - tmp_dir=user_working_directory, - ) - references = { - name: _task_output.OutputReference(_type_helpers.get_sdk_type_from_literal_type(variable.type)) - for name, variable in _six.iteritems(two_queries.interface.outputs) - } - - qubole_hive_jobs = two_queries._generate_plugin_objects(context, references) - assert len(qubole_hive_jobs) == 2 - - # deprecated, collection is only here for backwards compatibility - assert len(qubole_hive_jobs[0].query_collection.queries) == 1 - assert len(qubole_hive_jobs[1].query_collection.queries) == 1 - - # The output references should now have the same fake S3 path as the formatted queries - assert references["hive_results"].value[0].uri != "" - assert references["hive_results"].value[1].uri != "" - assert references["hive_results"].value[0].uri in qubole_hive_jobs[0].query.query - assert references["hive_results"].value[1].uri in qubole_hive_jobs[1].query.query - - -def test_hive_task_dynamic_job_spec_generation(): - with _common_utils.AutoDeletingTempDir("user_dir") as user_working_directory: - context = _common_engine.EngineContext( - execution_id=WorkflowExecutionIdentifier(project="unit_test", domain="unit_test", name="unit_test"), - execution_date=_datetime.utcnow(), - stats=None, # TODO: A mock stats object that we can read later. - logging=_logging, # TODO: A mock logging object that we can read later. - tmp_dir=user_working_directory, - ) - dj_spec = two_queries._produce_dynamic_job_spec(context, _literals.LiteralMap(literals={})) - - # Bindings - assert len(dj_spec.outputs[0].binding.collection.bindings) == 2 - assert isinstance(dj_spec.outputs[0].binding.collection.bindings[0].scalar.schema, Schema) - assert isinstance(dj_spec.outputs[0].binding.collection.bindings[1].scalar.schema, Schema) - - # Custom field is filled in - assert len(dj_spec.tasks[0].custom) > 0 diff --git a/tests/flytekit/unit/sdk/tasks/test_sidecar_tasks.py b/tests/flytekit/unit/sdk/tasks/test_sidecar_tasks.py deleted file mode 100644 index 10d1eb4fec..0000000000 --- a/tests/flytekit/unit/sdk/tasks/test_sidecar_tasks.py +++ /dev/null @@ -1,84 +0,0 @@ -import mock -from k8s.io.api.core.v1 import generated_pb2 - -from flytekit.common.tasks import sidecar_task as _sidecar_task -from flytekit.common.tasks import task as _sdk_task -from flytekit.configuration.internal import IMAGE as _IMAGE -from flytekit.models.core import identifier as _identifier -from flytekit.sdk.tasks import inputs, outputs, sidecar_task -from flytekit.sdk.types import Types - - -def get_pod_spec(): - a_container = generated_pb2.Container( - name="a container", - ) - a_container.command.extend(["fee", "fi", "fo", "fum"]) - a_container.volumeMounts.extend( - [ - generated_pb2.VolumeMount( - name="volume mount", - mountPath="some/where", - ) - ] - ) - - pod_spec = generated_pb2.PodSpec( - restartPolicy="OnFailure", - ) - pod_spec.containers.extend([a_container, generated_pb2.Container(name="another container")]) - return pod_spec - - -with mock.patch.object(_IMAGE, "get", return_value="docker.io/blah:abc123"): - - @inputs(in1=Types.Integer) - @outputs(out1=Types.String) - @sidecar_task( - cpu_request="10", - gpu_limit="2", - environment={"foo": "bar"}, - pod_spec=get_pod_spec(), - primary_container_name="a container", - annotations={"a": "a"}, - labels={"b": "b"}, - ) - def simple_sidecar_task(wf_params, in1, out1): - pass - - -simple_sidecar_task._id = _identifier.Identifier(_identifier.ResourceType.TASK, "project", "domain", "name", "version") - - -def test_sidecar_task(): - assert isinstance(simple_sidecar_task, _sdk_task.SdkTask) - assert isinstance(simple_sidecar_task, _sidecar_task.SdkSidecarTask) - - pod_spec = simple_sidecar_task.custom["podSpec"] - assert pod_spec["restartPolicy"] == "OnFailure" - assert len(pod_spec["containers"]) == 2 - primary_container = pod_spec["containers"][0] - assert primary_container["name"] == "a container" - assert primary_container["args"] == [ - "pyflyte-execute", - "--task-module", - "tests.flytekit.unit.sdk.tasks.test_sidecar_tasks", - "--task-name", - "simple_sidecar_task", - "--inputs", - "{{.input}}", - "--output-prefix", - "{{.outputPrefix}}", - "--raw-output-data-prefix", - "{{.rawOutputDataPrefix}}", - ] - assert primary_container["volumeMounts"] == [{"mountPath": "some/where", "name": "volume mount"}] - assert {"name": "foo", "value": "bar"} in primary_container["env"] - assert primary_container["resources"] == { - "requests": {"cpu": {"string": "10"}}, - "limits": {"gpu": {"string": "2"}}, - } - assert pod_spec["containers"][1]["name"] == "another container" - assert simple_sidecar_task.custom["primaryContainerName"] == "a container" - assert simple_sidecar_task.custom["annotations"]["a"] == "a" - assert simple_sidecar_task.custom["labels"]["b"] == "b" diff --git a/tests/flytekit/unit/sdk/tasks/test_spark_task.py b/tests/flytekit/unit/sdk/tasks/test_spark_task.py deleted file mode 100644 index cfccd5ecde..0000000000 --- a/tests/flytekit/unit/sdk/tasks/test_spark_task.py +++ /dev/null @@ -1,78 +0,0 @@ -import datetime as _datetime -import os as _os -import sys as _sys - -from flytekit.bin import entrypoint as _entrypoint -from flytekit.common import constants as _common_constants -from flytekit.common.tasks import sdk_runnable as _sdk_runnable -from flytekit.common.tasks import spark_task as _spark_task -from flytekit.models import types as _type_models -from flytekit.models.core import identifier as _identifier -from flytekit.sdk.tasks import inputs, outputs, spark_task -from flytekit.sdk.types import Types - - -@inputs(in1=Types.Integer) -@outputs(out1=Types.String) -@spark_task(spark_conf={"A": "B"}, hadoop_conf={"C": "D"}) -def default_task(wf_params, sc, in1, out1): - out1.set("hello") - - -default_task._id = _identifier.Identifier(_identifier.ResourceType.TASK, "project", "domain", "name", "version") - - -def test_default_python_task(): - assert isinstance(default_task, _spark_task.SdkSparkTask) - assert isinstance(default_task, _sdk_runnable.SdkRunnableTask) - assert default_task.interface.inputs["in1"].description == "" - assert default_task.interface.inputs["in1"].type == _type_models.LiteralType(simple=_type_models.SimpleType.INTEGER) - assert default_task.interface.outputs["out1"].description == "" - assert default_task.interface.outputs["out1"].type == _type_models.LiteralType( - simple=_type_models.SimpleType.STRING - ) - assert default_task.type == _common_constants.SdkTaskType.SPARK_TASK - assert default_task.task_function_name == "default_task" - assert default_task.task_module == __name__ - assert default_task.metadata.timeout == _datetime.timedelta(seconds=0) - assert default_task.metadata.deprecated_error_message == "" - assert default_task.metadata.discoverable is False - assert default_task.metadata.discovery_version == "" - assert default_task.metadata.retries.retries == 0 - assert len(default_task.container.resources.limits) == 0 - assert len(default_task.container.resources.requests) == 0 - assert default_task.custom["sparkConf"]["A"] == "B" - assert default_task.custom["hadoopConf"]["C"] == "D" - assert default_task.hadoop_conf["C"] == "D" - assert default_task.spark_conf["A"] == "B" - assert _os.path.abspath(_entrypoint.__file__)[:-1] in default_task.custom["mainApplicationFile"] - assert default_task.custom["executorPath"] == _sys.executable - - pb2 = default_task.to_flyte_idl() - assert pb2.custom["sparkConf"]["A"] == "B" - assert pb2.custom["hadoopConf"]["C"] == "D" - - -def test_overrides_spark_task(): - assert default_task.id.name == "name" - new_task = default_task.with_overrides(new_spark_conf={"x": "1"}, new_hadoop_conf={"y": "2"}) - assert isinstance(new_task, _spark_task.SdkSparkTask) - assert new_task.id.name.startswith("name-") - assert new_task.custom["sparkConf"]["x"] == "1" - assert new_task.custom["hadoopConf"]["y"] == "2" - - assert default_task.custom["sparkConf"]["A"] == "B" - assert default_task.custom["hadoopConf"]["C"] == "D" - - assert default_task.has_valid_name is False - default_task.assign_name("my-task") - assert default_task.has_valid_name - assert new_task.interface == default_task.interface - - assert default_task.__hash__() != new_task.__hash__() - - new_task2 = default_task.with_overrides(new_spark_conf={"x": "1"}, new_hadoop_conf={"y": "2"}) - assert new_task2.id.name == new_task.id.name - - t = new_task(in1=1) - assert t.outputs["out1"] is not None diff --git a/tests/flytekit/unit/sdk/tasks/test_tasks.py b/tests/flytekit/unit/sdk/tasks/test_tasks.py deleted file mode 100644 index 33e0287199..0000000000 --- a/tests/flytekit/unit/sdk/tasks/test_tasks.py +++ /dev/null @@ -1,108 +0,0 @@ -import datetime as _datetime -import os as _os - -from flytekit import configuration as _configuration -from flytekit.common import constants as _common_constants -from flytekit.common.tasks import sdk_runnable as _sdk_runnable -from flytekit.models import task as _task_models -from flytekit.models import types as _type_models -from flytekit.models.core import identifier as _identifier -from flytekit.sdk.tasks import inputs, outputs, python_task -from flytekit.sdk.types import Types - - -@inputs(in1=Types.Integer) -@outputs(out1=Types.String) -@python_task -def default_task(wf_params, in1, out1): - pass - - -default_task._id = _identifier.Identifier(_identifier.ResourceType.TASK, "project", "domain", "name", "version") - - -def test_default_python_task(): - assert isinstance(default_task, _sdk_runnable.SdkRunnableTask) - assert default_task.interface.inputs["in1"].description == "" - assert default_task.interface.inputs["in1"].type == _type_models.LiteralType(simple=_type_models.SimpleType.INTEGER) - assert default_task.interface.outputs["out1"].description == "" - assert default_task.interface.outputs["out1"].type == _type_models.LiteralType( - simple=_type_models.SimpleType.STRING - ) - assert default_task.type == _common_constants.SdkTaskType.PYTHON_TASK - assert default_task.task_function_name == "default_task" - assert default_task.task_module == __name__ - assert default_task.metadata.timeout == _datetime.timedelta(seconds=0) - assert default_task.metadata.deprecated_error_message == "" - assert default_task.metadata.discoverable is False - assert default_task.metadata.discovery_version == "" - assert default_task.metadata.retries.retries == 0 - assert len(default_task.container.resources.limits) == 0 - assert len(default_task.container.resources.requests) == 0 - - -def test_default_resources(): - with _configuration.TemporaryConfiguration( - _os.path.join( - _os.path.dirname(_os.path.realpath(__file__)), - "../../configuration/configs/good.config", - ) - ): - - @inputs(in1=Types.Integer) - @outputs(out1=Types.String) - @python_task() - def default_task2(wf_params, in1, out1): - pass - - request_map = {r.name: r.value for r in default_task2.container.resources.requests} - - limit_map = {l.name: l.value for l in default_task2.container.resources.limits} - - assert request_map[_task_models.Resources.ResourceName.CPU] == "500m" - assert request_map[_task_models.Resources.ResourceName.MEMORY] == "500Gi" - assert request_map[_task_models.Resources.ResourceName.GPU] == "1" - assert request_map[_task_models.Resources.ResourceName.STORAGE] == "500Gi" - - assert limit_map[_task_models.Resources.ResourceName.CPU] == "501m" - assert limit_map[_task_models.Resources.ResourceName.MEMORY] == "501Gi" - assert limit_map[_task_models.Resources.ResourceName.GPU] == "2" - assert limit_map[_task_models.Resources.ResourceName.STORAGE] == "501Gi" - - -def test_overriden_resources(): - with _configuration.TemporaryConfiguration( - _os.path.join( - _os.path.dirname(_os.path.realpath(__file__)), - "../../configuration/configs/good.config", - ) - ): - - @inputs(in1=Types.Integer) - @outputs(out1=Types.String) - @python_task( - memory_limit="100Gi", - memory_request="50Gi", - cpu_limit="1000m", - cpu_request="500m", - gpu_limit="1", - gpu_request="0", - storage_request="100Gi", - storage_limit="200Gi", - ) - def default_task2(wf_params, in1, out1): - pass - - request_map = {r.name: r.value for r in default_task2.container.resources.requests} - - limit_map = {l.name: l.value for l in default_task2.container.resources.limits} - - assert request_map[_task_models.Resources.ResourceName.CPU] == "500m" - assert request_map[_task_models.Resources.ResourceName.MEMORY] == "50Gi" - assert request_map[_task_models.Resources.ResourceName.GPU] == "0" - assert request_map[_task_models.Resources.ResourceName.STORAGE] == "100Gi" - - assert limit_map[_task_models.Resources.ResourceName.CPU] == "1000m" - assert limit_map[_task_models.Resources.ResourceName.MEMORY] == "100Gi" - assert limit_map[_task_models.Resources.ResourceName.GPU] == "1" - assert limit_map[_task_models.Resources.ResourceName.STORAGE] == "200Gi" diff --git a/tests/flytekit/unit/sdk/test_workflow.py b/tests/flytekit/unit/sdk/test_workflow.py deleted file mode 100644 index 2b604b5bf6..0000000000 --- a/tests/flytekit/unit/sdk/test_workflow.py +++ /dev/null @@ -1,150 +0,0 @@ -import pytest - -from flytekit.common import constants -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.types import base_sdk_types, containers, primitives -from flytekit.sdk.tasks import inputs, outputs, python_task -from flytekit.sdk.types import Types -from flytekit.sdk.workflow import Input, Output, workflow, workflow_class - - -def test_input(): - i = Input(primitives.Integer, help="blah", default=None) - assert i.name == "" - assert i.sdk_default is None - assert i.default == base_sdk_types.Void() - assert i.sdk_required is False - assert i.required is None - assert i.help == "blah" - assert i.var.description == "blah" - assert i.sdk_type == primitives.Integer - - i = i.rename_and_return_reference("new_name") - assert i.name == "new_name" - assert i.sdk_default is None - assert i.default == base_sdk_types.Void() - assert i.sdk_required is False - assert i.required is None - assert i.help == "blah" - assert i.var.description == "blah" - assert i.sdk_type == primitives.Integer - - i = Input(primitives.Integer, default=1) - assert i.name == "" - assert i.sdk_default == 1 - assert i.default == primitives.Integer(1) - assert i.sdk_required is False - assert i.required is None - assert i.help is None - assert i.var.description == "" - assert i.sdk_type == primitives.Integer - - i = i.rename_and_return_reference("new_name") - assert i.name == "new_name" - assert i.sdk_default == 1 - assert i.default == primitives.Integer(1) - assert i.sdk_required is False - assert i.required is None - assert i.help is None - assert i.var.description == "" - assert i.sdk_type == primitives.Integer - - with pytest.raises(_user_exceptions.FlyteAssertion): - Input(primitives.Integer, required=True, default=1) - - i = Input([primitives.Integer], default=[1, 2]) - assert i.name == "" - assert i.sdk_default == [1, 2] - assert i.default == containers.List(primitives.Integer)([primitives.Integer(1), primitives.Integer(2)]) - assert i.sdk_required is False - assert i.required is None - assert i.help is None - assert i.var.description == "" - assert i.sdk_type == containers.List(primitives.Integer) - - i = i.rename_and_return_reference("new_name") - assert i.name == "new_name" - assert i.sdk_default == [1, 2] - assert i.default == containers.List(primitives.Integer)([primitives.Integer(1), primitives.Integer(2)]) - assert i.sdk_required is False - assert i.required is None - assert i.help is None - assert i.var.description == "" - assert i.sdk_type == containers.List(primitives.Integer) - - -def test_output(): - o = Output(1, sdk_type=primitives.Integer, help="blah") - assert o.name == "" - assert o.var.description == "blah" - assert o.var.type == primitives.Integer.to_flyte_literal_type() - assert o.binding_data.scalar.primitive.integer == 1 - - o = o.rename_and_return_reference("new_name") - assert o.name == "new_name" - assert o.var.description == "blah" - assert o.var.type == primitives.Integer.to_flyte_literal_type() - assert o.binding_data.scalar.primitive.integer == 1 - - -def _get_node_by_id(wf, nid): - for n in wf.nodes: - if n.id == nid: - return n - assert False - - -def test_workflow_no_node_dependencies_or_outputs(): - @inputs(a=Types.Integer) - @outputs(b=Types.Integer) - @python_task - def my_task(wf_params, a, b): - b.set(a + 1) - - i1 = Input(Types.Integer) - i2 = Input(Types.Integer, default=5, help="Not required.") - - input_dict = {"input_1": i1, "input_2": i2} - - nodes = { - "a": my_task(a=input_dict["input_1"]), - "b": my_task(a=input_dict["input_2"]), - "c": my_task(a=100), - } - - w = workflow(inputs=input_dict, outputs={}, nodes=nodes) - - assert w.interface.inputs["input_1"].type == Types.Integer.to_flyte_literal_type() - assert w.interface.inputs["input_2"].type == Types.Integer.to_flyte_literal_type() - assert _get_node_by_id(w, "a").inputs[0].var == "a" - assert _get_node_by_id(w, "a").inputs[0].binding.promise.node_id == constants.GLOBAL_INPUT_NODE_ID - assert _get_node_by_id(w, "a").inputs[0].binding.promise.var == "input_1" - assert _get_node_by_id(w, "b").inputs[0].binding.promise.node_id == constants.GLOBAL_INPUT_NODE_ID - assert _get_node_by_id(w, "b").inputs[0].binding.promise.var == "input_2" - assert _get_node_by_id(w, "c").inputs[0].binding.scalar.primitive.integer == 100 - - -def test_workflow_metaclass_no_node_dependencies_or_outputs(): - @inputs(a=Types.Integer) - @outputs(b=Types.Integer) - @python_task - def my_task(wf_params, a, b): - b.set(a + 1) - - @workflow_class - class sup(object): - input_1 = Input(Types.Integer) - input_2 = Input(Types.Integer, default=5, help="Not required.") - - a = my_task(a=input_1) - b = my_task(a=input_2) - c = my_task(a=100) - - assert sup.interface.inputs["input_1"].type == Types.Integer.to_flyte_literal_type() - assert sup.interface.inputs["input_2"].type == Types.Integer.to_flyte_literal_type() - assert _get_node_by_id(sup, "a").inputs[0].var == "a" - assert _get_node_by_id(sup, "a").inputs[0].binding.promise.node_id == constants.GLOBAL_INPUT_NODE_ID - assert _get_node_by_id(sup, "a").inputs[0].binding.promise.var == "input_1" - assert _get_node_by_id(sup, "b").inputs[0].binding.promise.node_id == constants.GLOBAL_INPUT_NODE_ID - assert _get_node_by_id(sup, "b").inputs[0].binding.promise.var == "input_2" - assert _get_node_by_id(sup, "c").inputs[0].binding.scalar.primitive.integer == 100 diff --git a/tests/flytekit/unit/sdk/types/__init__.py b/tests/flytekit/unit/sdk/types/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/tests/flytekit/unit/sdk/types/test_blobs.py b/tests/flytekit/unit/sdk/types/test_blobs.py deleted file mode 100644 index 78e99fe5a7..0000000000 --- a/tests/flytekit/unit/sdk/types/test_blobs.py +++ /dev/null @@ -1,31 +0,0 @@ -import pytest - -from flytekit.common.types.impl import blobs as _blob_impl -from flytekit.sdk import types as _sdk_types - - -@pytest.mark.parametrize( - "blob_tuple", - [ - (_sdk_types.Types.Blob, _blob_impl.Blob), - (_sdk_types.Types.CSV, _blob_impl.Blob), - (_sdk_types.Types.MultiPartBlob, _blob_impl.MultiPartBlob), - (_sdk_types.Types.MultiPartCSV, _blob_impl.MultiPartBlob), - ], -) -def test_instantiable_blobs(blob_tuple): - sdk_type, impl = blob_tuple - - blob_inst = sdk_type() - blob_type_inst = sdk_type(blob_inst) - assert isinstance(blob_inst, impl) - assert isinstance(blob_type_inst, sdk_type) - - with pytest.raises(Exception): - sdk_type(1, 2) - - with pytest.raises(Exception): - sdk_type(a=1) - - blob_inst = sdk_type.create_at_known_location("abc") - assert isinstance(blob_inst, impl) diff --git a/tests/flytekit/unit/sdk/types/test_primitives.py b/tests/flytekit/unit/sdk/types/test_primitives.py deleted file mode 100644 index e039b9780b..0000000000 --- a/tests/flytekit/unit/sdk/types/test_primitives.py +++ /dev/null @@ -1,33 +0,0 @@ -import pytest - -from flytekit.sdk import types as _sdk_types - - -def test_integer(): - with pytest.raises(Exception): - _sdk_types.Types.Integer() - - -def test_float(): - with pytest.raises(Exception): - _sdk_types.Types.Float() - - -def test_string(): - with pytest.raises(Exception): - _sdk_types.Types.String() - - -def test_bool(): - with pytest.raises(Exception): - _sdk_types.Types.Boolean() - - -def test_datetime(): - with pytest.raises(Exception): - _sdk_types.Types.Datetime() - - -def test_timedelta(): - with pytest.raises(Exception): - _sdk_types.Types.Timedelta() diff --git a/tests/flytekit/unit/sdk/types/test_schema.py b/tests/flytekit/unit/sdk/types/test_schema.py deleted file mode 100644 index 95f70d6702..0000000000 --- a/tests/flytekit/unit/sdk/types/test_schema.py +++ /dev/null @@ -1,49 +0,0 @@ -import pytest - -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.sdk.tasks import inputs, outputs, python_task -from flytekit.sdk.types import Types - - -def test_generic_schema(): - @inputs(a=Types.Schema()) - @outputs(b=Types.Schema()) - @python_task - def fake_task(wf_params, a, b): - pass - - -def test_typed_schema(): - @inputs(a=Types.Schema([("a", Types.Integer), ("b", Types.Integer)])) - @outputs(b=Types.Schema([("a", Types.Integer), ("b", Types.Integer)])) - @python_task - def fake_task(wf_params, a, b): - pass - - -def test_bad_definition(): - with pytest.raises(_user_exceptions.FlyteValueException): - Types.Schema([]) - - -def test_bad_column_types(): - with pytest.raises(_user_exceptions.FlyteTypeException): - Types.Schema([("a", Types.Blob)]) - with pytest.raises(_user_exceptions.FlyteTypeException): - Types.Schema([("a", Types.MultiPartBlob)]) - with pytest.raises(_user_exceptions.FlyteTypeException): - Types.Schema([("a", Types.MultiPartCSV)]) - with pytest.raises(_user_exceptions.FlyteTypeException): - Types.Schema([("a", Types.CSV)]) - with pytest.raises(_user_exceptions.FlyteTypeException): - Types.Schema([("a", Types.Schema())]) - - -def test_create_from_hive_query(): - s, q = Types.Schema().create_from_hive_query("SELECT * FROM table", known_location="s3://somewhere/") - - assert s.mode == "wb" - assert s.local_path is None - assert s.remote_location == "s3://somewhere/" - assert "SELECT * FROM table" in q - assert s.remote_location in q diff --git a/tests/flytekit/unit/tasks/__init__.py b/tests/flytekit/unit/tasks/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/tests/flytekit/unit/test_plugins.py b/tests/flytekit/unit/test_plugins.py deleted file mode 100644 index c6680e34da..0000000000 --- a/tests/flytekit/unit/test_plugins.py +++ /dev/null @@ -1,48 +0,0 @@ -import pytest - -from flytekit import plugins -from flytekit.tools import lazy_loader - - -@pytest.mark.run(order=0) -def test_spark_plugin(): - plugins.pyspark.SparkContext - import pyspark - - assert plugins.pyspark.SparkContext == pyspark.SparkContext - - -@pytest.mark.run(order=1) -def test_schema_plugin(): - plugins.numpy.dtype - plugins.pandas.DataFrame - import numpy - import pandas - - assert plugins.numpy.dtype == numpy.dtype - assert pandas.DataFrame == pandas.DataFrame - - -@pytest.mark.run(order=2) -def test_sidecar_plugin(): - assert isinstance(plugins.k8s.io.api.core.v1.generated_pb2, lazy_loader._LazyLoadModule) - assert isinstance( - plugins.k8s.io.apimachinery.pkg.api.resource.generated_pb2, - lazy_loader._LazyLoadModule, - ) - import k8s.io.api.core.v1.generated_pb2 - import k8s.io.apimachinery.pkg.api.resource.generated_pb2 - - k8s.io.api.core.v1.generated_pb2.Container - k8s.io.apimachinery.pkg.api.resource.generated_pb2.Quantity - - -@pytest.mark.run(order=2) -def test_hive_sensor_plugin(): - assert isinstance(plugins.hmsclient, lazy_loader._LazyLoadModule) - assert isinstance(plugins.hmsclient.genthrift.hive_metastore.ttypes, lazy_loader._LazyLoadModule) - import hmsclient - import hmsclient.genthrift.hive_metastore.ttypes - - hmsclient.HMSClient - hmsclient.genthrift.hive_metastore.ttypes.NoSuchObjectException diff --git a/tests/flytekit/unit/common_tests/test_translator.py b/tests/flytekit/unit/test_translator.py similarity index 99% rename from tests/flytekit/unit/common_tests/test_translator.py rename to tests/flytekit/unit/test_translator.py index 91b1dd2780..e21353c38e 100644 --- a/tests/flytekit/unit/common_tests/test_translator.py +++ b/tests/flytekit/unit/test_translator.py @@ -2,7 +2,6 @@ from collections import OrderedDict from flytekit import ContainerTask, Resources -from flytekit.common.translator import get_serializable from flytekit.core import context_manager from flytekit.core.base_task import kwtypes from flytekit.core.context_manager import FastSerializationSettings, Image, ImageConfig @@ -11,6 +10,7 @@ from flytekit.core.task import ReferenceTask, task from flytekit.core.workflow import ReferenceWorkflow, workflow from flytekit.models.core import identifier as identifier_models +from flytekit.tools.translator import get_serializable default_img = Image(name="default", fqn="test", tag="tag") serialization_settings = context_manager.SerializationSettings( diff --git a/tests/flytekit/unit/tools/test_aws.py b/tests/flytekit/unit/tools/test_aws.py deleted file mode 100644 index 93445b0752..0000000000 --- a/tests/flytekit/unit/tools/test_aws.py +++ /dev/null @@ -1,7 +0,0 @@ -from flytekit.interfaces.data.s3.s3proxy import AwsS3Proxy - - -def test_aws_s3_splitting(): - (bucket, key) = AwsS3Proxy._split_s3_path_to_bucket_and_key("s3://bucket/some/key") - assert bucket == "bucket" - assert key == "some/key" diff --git a/tests/flytekit/unit/tools/test_lazy_loader.py b/tests/flytekit/unit/tools/test_lazy_loader.py deleted file mode 100644 index 7801318408..0000000000 --- a/tests/flytekit/unit/tools/test_lazy_loader.py +++ /dev/null @@ -1,14 +0,0 @@ -import pytest -import six - -from flytekit.tools import lazy_loader - - -def test_lazy_loader_error_message(): - lazy_mod = lazy_loader.lazy_load_module("made.up.module") - lazy_loader.LazyLoadPlugin("uninstalled_plugin", [], [lazy_mod]) - with pytest.raises(ImportError) as e: - lazy_mod.some_bad_attr - - assert "uninstalled_plugin" in six.text_type(e.value) - assert "flytekit[all]" in six.text_type(e.value) diff --git a/tests/flytekit/unit/tools/test_module_loader.py b/tests/flytekit/unit/tools/test_module_loader.py index aa7fdd255c..9a568f0260 100644 --- a/tests/flytekit/unit/tools/test_module_loader.py +++ b/tests/flytekit/unit/tools/test_module_loader.py @@ -1,12 +1,12 @@ import os import sys -from flytekit.common import utils as _utils +from flytekit.core import utils from flytekit.tools import module_loader def test_module_loading(): - with _utils.AutoDeletingTempDir("mypackage") as pkg: + with utils.AutoDeletingTempDir("mypackage") as pkg: path = pkg.name # Create directories top_level = os.path.join(path, "top") diff --git a/tests/flytekit/unit/type_engines/__init__.py b/tests/flytekit/unit/type_engines/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/tests/flytekit/unit/type_engines/default/__init__.py b/tests/flytekit/unit/type_engines/default/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/tests/flytekit/unit/type_engines/default/test_flyte_type_engine.py b/tests/flytekit/unit/type_engines/default/test_flyte_type_engine.py deleted file mode 100644 index 8157739456..0000000000 --- a/tests/flytekit/unit/type_engines/default/test_flyte_type_engine.py +++ /dev/null @@ -1,67 +0,0 @@ -import pytest -from flyteidl.core import errors_pb2 as _errors_pb2 - -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.common.types import proto as _proto -from flytekit.models import literals as _literal_models -from flytekit.models import types as _type_models -from flytekit.type_engines.default import flyte as _flyte_engine - - -def test_proto_from_literal_type(): - sdk_type = _flyte_engine.FlyteDefaultTypeEngine().get_sdk_type_from_literal_type( - _type_models.LiteralType( - simple=_type_models.SimpleType.BINARY, - metadata={_proto.Protobuf.PB_FIELD_KEY: "flyteidl.core.errors_pb2.ContainerError"}, - ) - ) - - assert sdk_type.pb_type == _errors_pb2.ContainerError - - -def test_generic_proto_from_literal_type(): - sdk_type = _flyte_engine.FlyteDefaultTypeEngine().get_sdk_type_from_literal_type( - _type_models.LiteralType( - simple=_type_models.SimpleType.STRUCT, - metadata={_proto.Protobuf.PB_FIELD_KEY: "flyteidl.core.errors_pb2.ContainerError"}, - ) - ) - - assert sdk_type.pb_type == _errors_pb2.ContainerError - - -def test_unloadable_module_from_literal_type(): - with pytest.raises(_user_exceptions.FlyteAssertion): - _flyte_engine.FlyteDefaultTypeEngine().get_sdk_type_from_literal_type( - _type_models.LiteralType( - simple=_type_models.SimpleType.BINARY, - metadata={_proto.Protobuf.PB_FIELD_KEY: "flyteidl.core.errors_pb2_no_exist.ContainerError"}, - ) - ) - - -def test_unloadable_proto_from_literal_type(): - with pytest.raises(_user_exceptions.FlyteAssertion): - _flyte_engine.FlyteDefaultTypeEngine().get_sdk_type_from_literal_type( - _type_models.LiteralType( - simple=_type_models.SimpleType.BINARY, - metadata={_proto.Protobuf.PB_FIELD_KEY: "flyteidl.core.errors_pb2.ContainerErrorNoExist"}, - ) - ) - - -def test_infer_proto_from_literal(): - sdk_type = _flyte_engine.FlyteDefaultTypeEngine().infer_sdk_type_from_literal( - _literal_models.Literal( - scalar=_literal_models.Scalar( - binary=_literal_models.Binary( - value="", - tag="{}{}".format( - _proto.Protobuf.TAG_PREFIX, - "flyteidl.core.errors_pb2.ContainerError", - ), - ) - ) - ) - ) - assert sdk_type.pb_type == _errors_pb2.ContainerError diff --git a/tests/flytekit/unit/use_scenarios/__init__.py b/tests/flytekit/unit/use_scenarios/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/tests/flytekit/unit/use_scenarios/unit_testing/__init__.py b/tests/flytekit/unit/use_scenarios/unit_testing/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/tests/flytekit/unit/use_scenarios/unit_testing/test_blobs.py b/tests/flytekit/unit/use_scenarios/unit_testing/test_blobs.py deleted file mode 100644 index 0530bdb865..0000000000 --- a/tests/flytekit/unit/use_scenarios/unit_testing/test_blobs.py +++ /dev/null @@ -1,165 +0,0 @@ -from flytekit.common.utils import AutoDeletingTempDir -from flytekit.sdk.tasks import inputs, outputs, python_task -from flytekit.sdk.test_utils import flyte_test -from flytekit.sdk.types import Types - - -@flyte_test -def test_create_blob_from_local_path(): - @outputs(a=Types.Blob) - @python_task - def test_create_from_local_path(wf_params, a): - with AutoDeletingTempDir("t") as tmp: - tmp_name = tmp.get_named_tempfile("abc.blob") - with open(tmp_name, "wb") as w: - w.write("Hello world".encode("utf-8")) - a.set(tmp_name) - - out = test_create_from_local_path.unit_test() - assert len(out) == 1 - with out["a"] as r: - assert r.read().decode("utf-8") == "Hello world" - - -@flyte_test -def test_write_blob(): - @outputs(a=Types.Blob) - @python_task - def test_write(wf_params, a): - b = Types.Blob() - with b as w: - w.write("Hello world".encode("utf-8")) - a.set(b) - - out = test_write.unit_test() - assert len(out) == 1 - with out["a"] as r: - assert r.read().decode("utf-8") == "Hello world" - - -@flyte_test -def test_blob_passing(): - @inputs(a=Types.Blob) - @outputs(b=Types.Blob) - @python_task - def test_pass(wf_params, a, b): - b.set(a) - - b = Types.Blob() - with b as w: - w.write("Hello world".encode("utf-8")) - - out = test_pass.unit_test(a=b) - assert len(out) == 1 - with out["b"] as r: - assert r.read().decode("utf-8") == "Hello world" - - out = test_pass.unit_test(a=out["b"]) - assert len(out) == 1 - with out["b"] as r: - assert r.read().decode("utf-8") == "Hello world" - - -@flyte_test -def test_create_multipartblob_from_local_path(): - @outputs(a=Types.MultiPartBlob) - @python_task - def test_create_from_local_path(wf_params, a): - with AutoDeletingTempDir("t") as tmp: - with open(tmp.get_named_tempfile("0"), "wb") as w: - w.write("Hello world".encode("utf-8")) - with open(tmp.get_named_tempfile("1"), "wb") as w: - w.write("Hello world2".encode("utf-8")) - a.set(tmp.name) - - out = test_create_from_local_path.unit_test() - assert len(out) == 1 - with out["a"] as r: - assert len(r) == 2 - assert r[0].read().decode("utf-8") == "Hello world" - assert r[1].read().decode("utf-8") == "Hello world2" - - -@flyte_test -def test_write_multipartblob(): - @outputs(a=Types.MultiPartBlob) - @python_task - def test_write(wf_params, a): - b = Types.MultiPartBlob() - with b.create_part("0") as w: - w.write("Hello world".encode("utf-8")) - with b.create_part("1") as w: - w.write("Hello world2".encode("utf-8")) - a.set(b) - - out = test_write.unit_test() - assert len(out) == 1 - with out["a"] as r: - assert len(r) == 2 - assert r[0].read().decode("utf-8") == "Hello world" - assert r[1].read().decode("utf-8") == "Hello world2" - - -@flyte_test -def test_multipartblob_passing(): - @inputs(a=Types.MultiPartBlob) - @outputs(b=Types.MultiPartBlob) - @python_task - def test_pass(wf_params, a, b): - b.set(a) - - b = Types.MultiPartBlob() - with b.create_part("0") as w: - w.write("Hello world".encode("utf-8")) - with b.create_part("1") as w: - w.write("Hello world2".encode("utf-8")) - - out = test_pass.unit_test(a=b) - assert len(out) == 1 - with out["b"] as r: - assert len(r) == 2 - assert r[0].read().decode("utf-8") == "Hello world" - assert r[1].read().decode("utf-8") == "Hello world2" - - out = test_pass.unit_test(a=out["b"]) - assert len(out) == 1 - with out["b"] as r: - assert len(r) == 2 - assert r[0].read().decode("utf-8") == "Hello world" - assert r[1].read().decode("utf-8") == "Hello world2" - - -@flyte_test -def test_write_csv(): - @outputs(a=Types.CSV) - @python_task - def test_write(wf_params, a): - b = Types.CSV() - with b as w: - w.write("Hello,world,hi") - a.set(b) - - out = test_write.unit_test() - assert len(out) == 1 - with out["a"] as r: - assert r.read() == "Hello,world,hi" - - -@flyte_test -def test_write_multipartcsv(): - @outputs(a=Types.MultiPartCSV) - @python_task - def test_write(wf_params, a): - b = Types.MultiPartCSV() - with b.create_part("0") as w: - w.write("Hello,world,1") - with b.create_part("1") as w: - w.write("Hello,world,2") - a.set(b) - - out = test_write.unit_test() - assert len(out) == 1 - with out["a"] as r: - assert len(r) == 2 - assert r[0].read() == "Hello,world,1" - assert r[1].read() == "Hello,world,2" diff --git a/tests/flytekit/unit/use_scenarios/unit_testing/test_hive_tasks.py b/tests/flytekit/unit/use_scenarios/unit_testing/test_hive_tasks.py deleted file mode 100644 index a137f254c1..0000000000 --- a/tests/flytekit/unit/use_scenarios/unit_testing/test_hive_tasks.py +++ /dev/null @@ -1,44 +0,0 @@ -import pytest - -from flytekit.sdk.tasks import hive_task - - -def test_no_queries(): - @hive_task - def test_hive_task(wf_params): - pass - - assert test_hive_task.unit_test() == [] - - -def test_empty_list_queries(): - @hive_task - def test_hive_task(wf_params): - return [] - - assert test_hive_task.unit_test() == [] - - -def test_one_query(): - @hive_task - def test_hive_task(wf_params): - return "abc" - - assert test_hive_task.unit_test() == ["abc"] - - -def test_multiple_queries(): - @hive_task - def test_hive_task(wf_params): - return ["abc", "cde"] - - assert test_hive_task.unit_test() == ["abc", "cde"] - - -def test_raise_exception(): - @hive_task - def test_hive_task(wf_params): - raise FloatingPointError("Floating point error for some reason.") - - with pytest.raises(FloatingPointError): - test_hive_task.unit_test() diff --git a/tests/flytekit/unit/use_scenarios/unit_testing/test_schemas.py b/tests/flytekit/unit/use_scenarios/unit_testing/test_schemas.py deleted file mode 100644 index e7df18e86a..0000000000 --- a/tests/flytekit/unit/use_scenarios/unit_testing/test_schemas.py +++ /dev/null @@ -1,139 +0,0 @@ -import pandas as pd -import pytest - -from flytekit.common.exceptions import user as _user_exceptions -from flytekit.sdk.tasks import inputs, outputs, python_task -from flytekit.sdk.test_utils import flyte_test -from flytekit.sdk.types import Types - - -@flyte_test -def test_generic_schema(): - @inputs(a=Types.Schema()) - @outputs(b=Types.Schema()) - @python_task - def copy_task(wf_params, a, b): - out = Types.Schema()() - with a as r: - with out as w: - for df in r.iter_chunks(): - w.write(df) - b.set(out) - - # Test generic copy and pass through - a = Types.Schema()() - with a as w: - w.write(pd.DataFrame.from_dict({"a": [1, 2, 3], "b": [4.0, 5.0, 6.0]})) - w.write(pd.DataFrame.from_dict({"a": [3, 2, 1], "b": [6.0, 5.0, 4.0]})) - - outs = copy_task.unit_test(a=a) - - with outs["b"] as r: - df = r.read() - assert list(df["a"]) == [1, 2, 3] - assert list(df["b"]) == [4.0, 5.0, 6.0] - - df = r.read() - assert list(df["a"]) == [3, 2, 1] - assert list(df["b"]) == [6.0, 5.0, 4.0] - - assert r.read() is None - - # Test typed copy and pass through - a = Types.Schema([("a", Types.Integer), ("b", Types.Float)])() - with a as w: - w.write(pd.DataFrame.from_dict({"a": [1, 2, 3], "b": [4.0, 5.0, 6.0]})) - w.write(pd.DataFrame.from_dict({"a": [3, 2, 1], "b": [6.0, 5.0, 4.0]})) - - outs = copy_task.unit_test(a=a) - - with outs["b"] as r: - df = r.read() - assert list(df["a"]) == [1, 2, 3] - assert list(df["b"]) == [4.0, 5.0, 6.0] - - df = r.read() - assert list(df["a"]) == [3, 2, 1] - assert list(df["b"]) == [6.0, 5.0, 4.0] - - assert r.read() is None - - -@flyte_test -def test_typed_schema(): - @inputs(a=Types.Schema([("a", Types.Integer), ("b", Types.Float)])) - @outputs(b=Types.Schema([("a", Types.Integer), ("b", Types.Float)])) - @python_task - def copy_task(wf_params, a, b): - out = Types.Schema([("a", Types.Integer), ("b", Types.Float)])() - with a as r: - with out as w: - for df in r.iter_chunks(): - w.write(df) - b.set(out) - - # Test typed copy and pass through - a = Types.Schema([("a", Types.Integer), ("b", Types.Float)])() - with a as w: - w.write(pd.DataFrame.from_dict({"a": [1, 2, 3], "b": [4.0, 5.0, 6.0]})) - w.write(pd.DataFrame.from_dict({"a": [3, 2, 1], "b": [6.0, 5.0, 4.0]})) - - outs = copy_task.unit_test(a=a) - - with outs["b"] as r: - df = r.read() - assert list(df["a"]) == [1, 2, 3] - assert list(df["b"]) == [4.0, 5.0, 6.0] - - df = r.read() - assert list(df["a"]) == [3, 2, 1] - assert list(df["b"]) == [6.0, 5.0, 4.0] - - assert r.read() is None - - # Test untyped failure - a = Types.Schema()() - with a as w: - w.write(pd.DataFrame.from_dict({"a": [1, 2, 3], "b": [4.0, 5.0, 6.0]})) - w.write(pd.DataFrame.from_dict({"a": [3, 2, 1], "b": [6.0, 5.0, 4.0]})) - - with pytest.raises(_user_exceptions.FlyteTypeException): - copy_task.unit_test(a=a) - - -@flyte_test -def test_subset_of_columns(): - @outputs(a=Types.Schema([("a", Types.Integer), ("b", Types.String)])) - @python_task() - def source(wf_params, a): - out = Types.Schema([("a", Types.Integer), ("b", Types.String)])() - with out as writer: - writer.write(pd.DataFrame.from_dict({"a": [1, 2, 3, 4, 5], "b": ["a", "b", "c", "d", "e"]})) - a.set(out) - - @inputs(a=Types.Schema([("a", Types.Integer)])) - @python_task() - def sink(wf_params, a): - with a as reader: - df = reader.read(concat=True) - assert len(df.columns.values) == 1 - assert df["a"].tolist() == [1, 2, 3, 4, 5] - - with a as reader: - df = reader.read(truncate_extra_columns=False) - assert df.columns.values.tolist() == ["a", "b"] - assert df["a"].tolist() == [1, 2, 3, 4, 5] - assert df["b"].tolist() == ["a", "b", "c", "d", "e"] - - o = source.unit_test() - sink.unit_test(**o) - - -@flyte_test -def test_no_output_set(): - @outputs(a=Types.Schema()) - @python_task() - def null_set(wf_params, a): - pass - - assert null_set.unit_test()["a"] is None From a50f99155a2fd5c3077f5d9ea4e6f50de874d5d8 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Tue, 18 Jan 2022 01:26:10 +0800 Subject: [PATCH 056/128] Add Bigquery plugin (#789) * Add bigquery plugin Signed-off-by: Kevin Su * Update dependency Signed-off-by: Kevin Su * update get_custom Signed-off-by: Kevin Su * Add structured dataset Signed-off-by: Kevin Su * Add structured dataset Signed-off-by: Kevin Su * Updated comment Signed-off-by: Kevin Su * Add BQ in GA Signed-off-by: Kevin Su * alphabetical order Signed-off-by: Kevin Su Signed-off-by: maximsmol --- .github/workflows/pythonbuild.yml | 8 +- plugins/flytekit-bigquery/README.md | 11 + .../flytekitplugins/bigquery/__init__.py | 1 + .../flytekitplugins/bigquery/task.py | 82 +++++++ plugins/flytekit-bigquery/requirements.in | 2 + plugins/flytekit-bigquery/requirements.txt | 202 ++++++++++++++++++ plugins/flytekit-bigquery/setup.py | 36 ++++ plugins/flytekit-bigquery/tests/__init__.py | 0 .../flytekit-bigquery/tests/test_bigquery.py | 71 ++++++ plugins/setup.py | 5 +- 10 files changed, 414 insertions(+), 4 deletions(-) create mode 100644 plugins/flytekit-bigquery/README.md create mode 100644 plugins/flytekit-bigquery/flytekitplugins/bigquery/__init__.py create mode 100644 plugins/flytekit-bigquery/flytekitplugins/bigquery/task.py create mode 100644 plugins/flytekit-bigquery/requirements.in create mode 100644 plugins/flytekit-bigquery/requirements.txt create mode 100644 plugins/flytekit-bigquery/setup.py create mode 100644 plugins/flytekit-bigquery/tests/__init__.py create mode 100644 plugins/flytekit-bigquery/tests/test_bigquery.py diff --git a/.github/workflows/pythonbuild.yml b/.github/workflows/pythonbuild.yml index 59085ec468..0ef81526fe 100644 --- a/.github/workflows/pythonbuild.yml +++ b/.github/workflows/pythonbuild.yml @@ -61,8 +61,10 @@ jobs: matrix: python-version: ["3.8", "3.9", "3.10"] plugin-names: + # Please maintain an alphabetical order in the following list - flytekit-aws-athena - flytekit-aws-sagemaker + - flytekit-bigquery - flytekit-data-fsspec - flytekit-dolt - flytekit-greatexpectations @@ -71,12 +73,12 @@ jobs: - flytekit-kf-mpi - flytekit-kf-pytorch - flytekit-kf-tensorflow + - flytekit-modin + - flytekit-pandera - flytekit-papermill + - flytekit-snowflake - flytekit-spark - flytekit-sqlalchemy - - flytekit-pandera - - flytekit-snowflake - - flytekit-modin exclude: # flytekit-modin depends on ray which does not have a 3.10 wheel yet. # Issue tracked in https://github.com/ray-project/ray/issues/19116. diff --git a/plugins/flytekit-bigquery/README.md b/plugins/flytekit-bigquery/README.md new file mode 100644 index 0000000000..7b8468ffc2 --- /dev/null +++ b/plugins/flytekit-bigquery/README.md @@ -0,0 +1,11 @@ +# Flytekit BigQuery Plugin + +BigQuery enables us to build data-intensive applications without operational burden. Flyte backend can be connected with the BigQuery service. Once enabled, it can allow you to query a BigQuery table. + +To install the plugin, run the following command: + +```bash +pip install flytekitplugins-bigquery +``` + +To configure BigQuery in the Flyte deployment's backend, follow the [configuration guide](https://docs.flyte.org/en/latest/deployment/plugin_setup/gcp/bigquery.html#deployment-plugin-setup-gcp-bigquery). diff --git a/plugins/flytekit-bigquery/flytekitplugins/bigquery/__init__.py b/plugins/flytekit-bigquery/flytekitplugins/bigquery/__init__.py new file mode 100644 index 0000000000..cb259e4f49 --- /dev/null +++ b/plugins/flytekit-bigquery/flytekitplugins/bigquery/__init__.py @@ -0,0 +1 @@ +from .task import BigQueryConfig, BigQueryTask diff --git a/plugins/flytekit-bigquery/flytekitplugins/bigquery/task.py b/plugins/flytekit-bigquery/flytekitplugins/bigquery/task.py new file mode 100644 index 0000000000..b7a5104dea --- /dev/null +++ b/plugins/flytekit-bigquery/flytekitplugins/bigquery/task.py @@ -0,0 +1,82 @@ +from dataclasses import dataclass +from typing import Any, Dict, Optional, Type + +from google.cloud import bigquery +from google.protobuf import json_format +from google.protobuf.struct_pb2 import Struct + +from flytekit import StructuredDataset +from flytekit.extend import SerializationSettings, SQLTask +from flytekit.models import task as _task_model + + +@dataclass +class BigQueryConfig(object): + """ + BigQueryConfig should be used to configure a BigQuery Task. + """ + + ProjectID: str + Location: Optional[str] = None + QueryJobConfig: Optional[bigquery.QueryJobConfig] = None + + +class BigQueryTask(SQLTask[BigQueryConfig]): + """ + This is the simplest form of a BigQuery Task, that can be used even for tasks that do not produce any output. + """ + + # This task is executed using the BigQuery handler in the backend. + # https://github.com/flyteorg/flyteplugins/blob/43623826fb189fa64dc4cb53e7025b517d911f22/go/tasks/plugins/webapi/bigquery/plugin.go#L34 + _TASK_TYPE = "bigquery_query_job_task" + + def __init__( + self, + name: str, + query_template: str, + task_config: Optional[BigQueryConfig], + inputs: Optional[Dict[str, Type]] = None, + output_structured_dataset_type: Optional[Type[StructuredDataset]] = None, + **kwargs, + ): + """ + To be used to query BigQuery Tables. + + :param name: Name of this task, should be unique in the project + :param query_template: The actual query to run. We use Flyte's Golang templating format for Query templating. + Refer to the templating documentation + :param task_config: BigQueryConfig object + :param inputs: Name and type of inputs specified as an ordered dictionary + :param output_structured_dataset_type: If some data is produced by this query, then you can specify the output StructuredDataset type + :param kwargs: All other args required by Parent type - SQLTask + """ + outputs = None + if output_structured_dataset_type is not None: + outputs = { + "results": output_structured_dataset_type, + } + super().__init__( + name=name, + task_config=task_config, + query_template=query_template, + inputs=inputs, + outputs=outputs, + task_type=self._TASK_TYPE, + **kwargs, + ) + self._output_structured_dataset_type = output_structured_dataset_type + + def get_custom(self, settings: SerializationSettings) -> Dict[str, Any]: + config = { + "Location": self.task_config.Location, + "ProjectID": self.task_config.ProjectID, + } + if self.task_config.QueryJobConfig is not None: + config.update(self.task_config.QueryJobConfig.to_api_repr()["query"]) + s = Struct() + s.update(config) + return json_format.MessageToDict(s) + + def get_sql(self, settings: SerializationSettings) -> Optional[_task_model.Sql]: + sql = _task_model.Sql(statement=self.query_template, dialect=_task_model.Sql.Dialect.ANSI) + return sql diff --git a/plugins/flytekit-bigquery/requirements.in b/plugins/flytekit-bigquery/requirements.in new file mode 100644 index 0000000000..d01f21ac2b --- /dev/null +++ b/plugins/flytekit-bigquery/requirements.in @@ -0,0 +1,2 @@ +. +-e file:.#egg=flytekitplugins-bigquery diff --git a/plugins/flytekit-bigquery/requirements.txt b/plugins/flytekit-bigquery/requirements.txt new file mode 100644 index 0000000000..45be4f75a0 --- /dev/null +++ b/plugins/flytekit-bigquery/requirements.txt @@ -0,0 +1,202 @@ +# +# This file is autogenerated by pip-compile with python 3.9 +# To update, run: +# +# pip-compile requirements.in +# +-e file:.#egg=flytekitplugins-bigquery + # via -r requirements.in +arrow==1.2.1 + # via jinja2-time +binaryornot==0.4.4 + # via cookiecutter +cachetools==4.2.4 + # via google-auth +certifi==2021.10.8 + # via requests +chardet==4.0.0 + # via binaryornot +charset-normalizer==2.0.7 + # via requests +checksumdir==1.2.0 + # via flytekit +click==7.1.2 + # via + # cookiecutter + # flytekit +cloudpickle==2.0.0 + # via flytekit +cookiecutter==1.7.3 + # via flytekit +croniter==1.0.15 + # via flytekit +dataclasses-json==0.5.6 + # via flytekit +decorator==5.1.0 + # via retry +deprecated==1.2.13 + # via flytekit +diskcache==5.2.1 + # via flytekit +docker-image-py==0.1.12 + # via flytekit +docstring-parser==0.12 + # via flytekit +flyteidl==0.21.8 + # via flytekit +flytekit==0.24.0 + # via flytekitplugins-bigquery +google-api-core[grpc]==2.3.2 + # via + # google-cloud-bigquery + # google-cloud-core +google-auth==2.3.3 + # via + # google-api-core + # google-cloud-core +google-cloud-bigquery==2.31.0 + # via flytekitplugins-bigquery +google-cloud-core==2.2.1 + # via google-cloud-bigquery +google-crc32c==1.3.0 + # via google-resumable-media +google-resumable-media==2.1.0 + # via google-cloud-bigquery +googleapis-common-protos==1.54.0 + # via + # google-api-core + # grpcio-status +grpcio==1.43.0 + # via + # flytekit + # google-api-core + # google-cloud-bigquery + # grpcio-status +grpcio-status==1.43.0 + # via google-api-core +idna==3.3 + # via requests +importlib-metadata==4.8.2 + # via keyring +jinja2==3.0.3 + # via + # cookiecutter + # jinja2-time +jinja2-time==0.2.0 + # via cookiecutter +keyring==23.2.1 + # via flytekit +markupsafe==2.0.1 + # via jinja2 +marshmallow==3.14.0 + # via + # dataclasses-json + # marshmallow-enum + # marshmallow-jsonschema +marshmallow-enum==1.5.1 + # via dataclasses-json +marshmallow-jsonschema==0.13.0 + # via flytekit +mypy-extensions==0.4.3 + # via typing-inspect +natsort==8.0.0 + # via flytekit +numpy==1.21.4 + # via + # pandas + # pyarrow +packaging==21.3 + # via google-cloud-bigquery +pandas==1.3.4 + # via flytekit +poyo==0.5.0 + # via cookiecutter +proto-plus==1.19.8 + # via google-cloud-bigquery +protobuf==3.19.1 + # via + # flyteidl + # flytekit + # google-api-core + # google-cloud-bigquery + # googleapis-common-protos + # grpcio-status + # proto-plus +py==1.11.0 + # via retry +pyarrow==6.0.0 + # via flytekit +pyasn1==0.4.8 + # via + # pyasn1-modules + # rsa +pyasn1-modules==0.2.8 + # via google-auth +pyparsing==3.0.6 + # via packaging +python-dateutil==2.8.1 + # via + # arrow + # croniter + # flytekit + # google-cloud-bigquery + # pandas +python-json-logger==2.0.2 + # via flytekit +python-slugify==5.0.2 + # via cookiecutter +pytimeparse==1.1.8 + # via flytekit +pytz==2018.4 + # via + # flytekit + # pandas +regex==2021.11.10 + # via docker-image-py +requests==2.26.0 + # via + # cookiecutter + # flytekit + # google-api-core + # google-cloud-bigquery + # responses +responses==0.15.0 + # via flytekit +retry==0.9.2 + # via flytekit +rsa==4.8 + # via google-auth +six==1.16.0 + # via + # cookiecutter + # flytekit + # google-auth + # grpcio + # python-dateutil + # responses +sortedcontainers==2.4.0 + # via flytekit +statsd==3.3.0 + # via flytekit +text-unidecode==1.3 + # via python-slugify +typing-extensions==3.10.0.2 + # via typing-inspect +typing-inspect==0.7.1 + # via dataclasses-json +urllib3==1.26.7 + # via + # flytekit + # requests + # responses +wheel==0.37.0 + # via flytekit +wrapt==1.13.3 + # via + # deprecated + # flytekit +zipp==3.6.0 + # via importlib-metadata + +# The following packages are considered to be unsafe in a requirements file: +# setuptools diff --git a/plugins/flytekit-bigquery/setup.py b/plugins/flytekit-bigquery/setup.py new file mode 100644 index 0000000000..7be6fc7d56 --- /dev/null +++ b/plugins/flytekit-bigquery/setup.py @@ -0,0 +1,36 @@ +from setuptools import setup + +PLUGIN_NAME = "bigquery" + +microlib_name = f"flytekitplugins-{PLUGIN_NAME}" + +plugin_requires = ["flytekit>=v0.30.0b3,<1.0.0", "google-cloud-bigquery"] + +__version__ = "0.0.0+develop" + +setup( + name=microlib_name, + version=__version__, + author="flyteorg", + author_email="admin@flyte.org", + description="This package holds the Bigquery plugins for flytekit", + namespace_packages=["flytekitplugins"], + packages=[f"flytekitplugins.{PLUGIN_NAME}"], + install_requires=plugin_requires, + license="apache2", + python_requires=">=3.7", + classifiers=[ + "Intended Audience :: Science/Research", + "Intended Audience :: Developers", + "License :: OSI Approved :: Apache Software License", + "Programming Language :: Python :: 3.7", + "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Topic :: Scientific/Engineering", + "Topic :: Scientific/Engineering :: Artificial Intelligence", + "Topic :: Software Development", + "Topic :: Software Development :: Libraries", + "Topic :: Software Development :: Libraries :: Python Modules", + ], +) diff --git a/plugins/flytekit-bigquery/tests/__init__.py b/plugins/flytekit-bigquery/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/plugins/flytekit-bigquery/tests/test_bigquery.py b/plugins/flytekit-bigquery/tests/test_bigquery.py new file mode 100644 index 0000000000..78d6c0893f --- /dev/null +++ b/plugins/flytekit-bigquery/tests/test_bigquery.py @@ -0,0 +1,71 @@ +from collections import OrderedDict + +import pytest +from flytekitplugins.bigquery import BigQueryConfig, BigQueryTask +from google.cloud.bigquery import QueryJobConfig +from google.protobuf import json_format +from google.protobuf.struct_pb2 import Struct + +from flytekit import StructuredDataset, kwtypes, workflow +from flytekit.extend import Image, ImageConfig, SerializationSettings, get_serializable + +query_template = "SELECT * FROM `bigquery-public-data.crypto_dogecoin.transactions` WHERE @version = 1 LIMIT 10" + + +def test_serialization(): + bigquery_task = BigQueryTask( + name="flytekit.demo.bigquery_task.query", + inputs=kwtypes(ds=str), + task_config=BigQueryConfig( + ProjectID="Flyte", Location="Asia", QueryJobConfig=QueryJobConfig(allow_large_results=True) + ), + query_template=query_template, + output_structured_dataset_type=StructuredDataset, + ) + + @workflow + def my_wf(ds: str) -> StructuredDataset: + return bigquery_task(ds=ds) + + default_img = Image(name="default", fqn="test", tag="tag") + serialization_settings = SerializationSettings( + project="proj", + domain="dom", + version="123", + image_config=ImageConfig(default_image=default_img, images=[default_img]), + env={}, + ) + + task_spec = get_serializable(OrderedDict(), serialization_settings, bigquery_task) + + assert "SELECT * FROM `bigquery-public-data.crypto_dogecoin.transactions`" in task_spec.template.sql.statement + assert "@version" in task_spec.template.sql.statement + assert task_spec.template.sql.dialect == task_spec.template.sql.Dialect.ANSI + s = Struct() + s.update({"ProjectID": "Flyte", "Location": "Asia", "allowLargeResults": True}) + assert task_spec.template.custom == json_format.MessageToDict(s) + assert len(task_spec.template.interface.inputs) == 1 + assert len(task_spec.template.interface.outputs) == 1 + + admin_workflow_spec = get_serializable(OrderedDict(), serialization_settings, my_wf) + assert admin_workflow_spec.template.interface.outputs["o0"].type.structured_dataset_type is not None + assert admin_workflow_spec.template.outputs[0].var == "o0" + assert admin_workflow_spec.template.outputs[0].binding.promise.node_id == "n0" + assert admin_workflow_spec.template.outputs[0].binding.promise.var == "results" + + +def test_local_exec(): + bigquery_task = BigQueryTask( + name="flytekit.demo.bigquery_task.query2", + inputs=kwtypes(ds=str), + query_template=query_template, + task_config=BigQueryConfig(ProjectID="Flyte", Location="Asia"), + output_structured_dataset_type=StructuredDataset, + ) + + assert len(bigquery_task.interface.inputs) == 1 + assert len(bigquery_task.interface.outputs) == 1 + + # will not run locally + with pytest.raises(Exception): + bigquery_task() diff --git a/plugins/setup.py b/plugins/setup.py index 4c5b69dc36..8f3cc5c299 100644 --- a/plugins/setup.py +++ b/plugins/setup.py @@ -6,9 +6,12 @@ from setuptools.command.install import install PACKAGE_NAME = "flytekitplugins-parent" + +# Please maintain an alphabetical order in the following list SOURCES = { "flytekitplugins-athena": "flytekit-aws-athena", "flytekitplugins-awssagemaker": "flytekit-aws-sagemaker", + "flytekitplugins-bigquery": "flytekit-bigquery", "flytekitplugins-fsspec": "flytekit-data-fsspec", "flytekitplugins-dolt": "flytekit-dolt", "flytekitplugins-great_expectations": "flytekit-greatexpectations", @@ -17,12 +20,12 @@ "flytekitplugins-kfmpi": "flytekit-kf-mpi", "flytekitplugins-kfpytorch": "flytekit-kf-pytorch", "flytekitplugins-kftensorflow": "flytekit-kf-tensorflow", + "flytekitplugins-modin": "flytekit-modin", "flytekitplugins-pandera": "flytekit-pandera", "flytekitplugins-papermill": "flytekit-papermill", "flytekitplugins-snowflake": "flytekit-snowflake", "flytekitplugins-spark": "flytekit-spark", "flytekitplugins-sqlalchemy": "flytekit-sqlalchemy", - "flytekitplugins-modin": "flytekit-modin", } From 1220976097da38b19e686626f7eae921574c22f3 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Tue, 18 Jan 2022 03:00:16 +0800 Subject: [PATCH 057/128] Fixed flytekit-papermill ImportError (#818) Signed-off-by: Kevin Su --- .github/workflows/pythonbuild.yml | 2 +- .../flytekit-papermill/dev-requirements.in | 1 + .../flytekit-papermill/dev-requirements.txt | 152 ++++++++++++++++++ plugins/flytekit-papermill/requirements.in | 1 - plugins/flytekit-papermill/requirements.txt | 30 +--- plugins/flytekit-papermill/setup.py | 1 - 6 files changed, 160 insertions(+), 27 deletions(-) create mode 100644 plugins/flytekit-papermill/dev-requirements.in create mode 100644 plugins/flytekit-papermill/dev-requirements.txt diff --git a/.github/workflows/pythonbuild.yml b/.github/workflows/pythonbuild.yml index 0ef81526fe..2eeabd109d 100644 --- a/.github/workflows/pythonbuild.yml +++ b/.github/workflows/pythonbuild.yml @@ -103,8 +103,8 @@ jobs: make setup cd plugins/${{ matrix.plugin-names }} pip install -e . + if [ -f dev-requirements.txt ]; then pip install -r dev-requirements.txt; fi pip install --no-deps -U https://github.com/flyteorg/flytekit/archive/${{ github.sha }}.zip#egg=flytekit - git clone https://github.com/flyteorg/flyteidl.git && cd flyteidl && git checkout union_type && pip install . && cd .. pip freeze - name: Test with coverage run: | diff --git a/plugins/flytekit-papermill/dev-requirements.in b/plugins/flytekit-papermill/dev-requirements.in new file mode 100644 index 0000000000..8bf53e2636 --- /dev/null +++ b/plugins/flytekit-papermill/dev-requirements.in @@ -0,0 +1 @@ +flytekitplugins-spark>=0.30.0b4 diff --git a/plugins/flytekit-papermill/dev-requirements.txt b/plugins/flytekit-papermill/dev-requirements.txt new file mode 100644 index 0000000000..849cef4c9f --- /dev/null +++ b/plugins/flytekit-papermill/dev-requirements.txt @@ -0,0 +1,152 @@ +# +# This file is autogenerated by pip-compile with python 3.9 +# To update, run: +# +# pip-compile dev-requirements.in +# +arrow==1.2.1 + # via jinja2-time +binaryornot==0.4.4 + # via cookiecutter +certifi==2021.10.8 + # via requests +chardet==4.0.0 + # via binaryornot +charset-normalizer==2.0.10 + # via requests +checksumdir==1.2.0 + # via flytekit +click==7.1.2 + # via + # cookiecutter + # flytekit +cloudpickle==2.0.0 + # via flytekit +cookiecutter==1.7.3 + # via flytekit +croniter==1.2.0 + # via flytekit +dataclasses-json==0.5.6 + # via flytekit +decorator==5.1.1 + # via retry +deprecated==1.2.13 + # via flytekit +diskcache==5.4.0 + # via flytekit +docker-image-py==0.1.12 + # via flytekit +docstring-parser==0.13 + # via flytekit +flyteidl==0.21.23 + # via flytekit +flytekit==0.26.0 + # via flytekitplugins-spark +flytekitplugins-spark==0.30.0b4 + # via -r dev-requirements.in +grpcio==1.43.0 + # via flytekit +idna==3.3 + # via requests +importlib-metadata==4.10.1 + # via keyring +jinja2==3.0.3 + # via + # cookiecutter + # jinja2-time +jinja2-time==0.2.0 + # via cookiecutter +keyring==23.5.0 + # via flytekit +markupsafe==2.0.1 + # via jinja2 +marshmallow==3.14.1 + # via + # dataclasses-json + # marshmallow-enum + # marshmallow-jsonschema +marshmallow-enum==1.5.1 + # via dataclasses-json +marshmallow-jsonschema==0.13.0 + # via flytekit +mypy-extensions==0.4.3 + # via typing-inspect +natsort==8.0.2 + # via flytekit +numpy==1.22.1 + # via + # pandas + # pyarrow +pandas==1.3.5 + # via flytekit +poyo==0.5.0 + # via cookiecutter +protobuf==3.19.3 + # via + # flyteidl + # flytekit +py==1.11.0 + # via retry +py4j==0.10.9.2 + # via pyspark +pyarrow==6.0.1 + # via flytekit +pyspark==3.2.0 + # via flytekitplugins-spark +python-dateutil==2.8.1 + # via + # arrow + # croniter + # flytekit + # pandas +python-json-logger==2.0.2 + # via flytekit +python-slugify==5.0.2 + # via cookiecutter +pytimeparse==1.1.8 + # via flytekit +pytz==2021.3 + # via + # flytekit + # pandas +regex==2021.11.10 + # via docker-image-py +requests==2.27.1 + # via + # cookiecutter + # flytekit + # responses +responses==0.17.0 + # via flytekit +retry==0.9.2 + # via flytekit +six==1.16.0 + # via + # cookiecutter + # flytekit + # grpcio + # python-dateutil + # responses +sortedcontainers==2.4.0 + # via flytekit +statsd==3.3.0 + # via flytekit +text-unidecode==1.3 + # via python-slugify +typing-extensions==4.0.1 + # via typing-inspect +typing-inspect==0.7.1 + # via dataclasses-json +urllib3==1.26.8 + # via + # flytekit + # requests + # responses +wheel==0.37.1 + # via flytekit +wrapt==1.13.3 + # via + # deprecated + # flytekit +zipp==3.7.0 + # via importlib-metadata diff --git a/plugins/flytekit-papermill/requirements.in b/plugins/flytekit-papermill/requirements.in index 668a9fe9fa..a9dab93776 100644 --- a/plugins/flytekit-papermill/requirements.in +++ b/plugins/flytekit-papermill/requirements.in @@ -1,3 +1,2 @@ . -e file:.#egg=flytekitplugins-papermill -flytekitplugins-spark diff --git a/plugins/flytekit-papermill/requirements.txt b/plugins/flytekit-papermill/requirements.txt index c31e0824bf..9c5a82d050 100644 --- a/plugins/flytekit-papermill/requirements.txt +++ b/plugins/flytekit-papermill/requirements.txt @@ -8,6 +8,10 @@ # via -r requirements.in ansiwrap==0.8.4 # via papermill +appnope==0.1.2 + # via + # ipykernel + # ipython arrow==1.2.1 # via jinja2-time attrs==21.2.0 @@ -22,8 +26,6 @@ bleach==4.1.0 # via nbconvert certifi==2021.10.8 # via requests -cffi==1.15.0 - # via cryptography chardet==4.0.0 # via binaryornot charset-normalizer==2.0.7 @@ -42,8 +44,6 @@ cookiecutter==1.7.3 # via flytekit croniter==1.0.15 # via flytekit -cryptography==35.0.0 - # via secretstorage dataclasses-json==0.5.6 # via flytekit debugpy==1.5.1 @@ -67,16 +67,10 @@ entrypoints==0.3 # jupyter-client # nbconvert # papermill -flyteidl==0.21.8 +flyteidl==0.21.23 # via flytekit flytekit==0.24.0 - # via - # flytekitplugins-papermill - # flytekitplugins-spark -flytekitplugins-spark==0.24.0 - # via - # -r requirements.in - # flytekitplugins-papermill + # via flytekitplugins-papermill grpcio==1.41.1 # via flytekit idna==3.3 @@ -91,10 +85,6 @@ ipython-genutils==0.2.0 # via nbformat jedi==0.18.0 # via ipython -jeepney==0.7.1 - # via - # keyring - # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -189,12 +179,8 @@ ptyprocess==0.7.0 # via pexpect py==1.11.0 # via retry -py4j==0.10.9.2 - # via pyspark pyarrow==6.0.0 # via flytekit -pycparser==2.21 - # via cffi pygments==2.10.0 # via # ipython @@ -204,8 +190,6 @@ pyparsing==2.4.7 # via packaging pyrsistent==0.18.0 # via jsonschema -pyspark==3.2.0 - # via flytekitplugins-spark python-dateutil==2.8.1 # via # arrow @@ -241,8 +225,6 @@ responses==0.15.0 # via flytekit retry==0.9.2 # via flytekit -secretstorage==3.3.1 - # via keyring six==1.16.0 # via # bleach diff --git a/plugins/flytekit-papermill/setup.py b/plugins/flytekit-papermill/setup.py index 4fa4943302..f8e92c92d0 100644 --- a/plugins/flytekit-papermill/setup.py +++ b/plugins/flytekit-papermill/setup.py @@ -6,7 +6,6 @@ plugin_requires = [ "flytekit>=0.16.0b0,<1.0.0", - "flytekitplugins-spark>=0.16.0b0,<1.0.0,!=0.24.0b0", "papermill>=1.2.0", "nbconvert>=6.0.7", "ipykernel>=5.0.0", From bcc4f61e3a45076eb89a6a8c4ebdd3d1fb4bc19c Mon Sep 17 00:00:00 2001 From: Zach Palchick Date: Mon, 17 Jan 2022 20:53:59 -0800 Subject: [PATCH 058/128] Add support for string-format-like sytax for shell task (#792) * POC: Add support for f-string like sytax for shell task This commit is a proof of concept adding f-string like syntax for shell_tasks. This supports using nested types for script inputs, such as data classes. This change was motivated by the desire to combine shell_tasks that have multiple inputs with map_tasks which only support tasks with a single input. This commit is only a starting point, since it makes some changes to the shell_task API (adds a template_style field), and modifies some of the default behavior for ease of implementation (e.g. throwing an error when there are unused input arguments). Signed-off-by: Zach Palchick * Drop support for old/regex style for doing string interpolation Signed-off-by: Zach Palchick Signed-off-by: maximsmol --- flytekit/extras/tasks/shell.py | 96 ++++++++-------- .../flytekit/unit/extras/tasks/test_shell.py | 106 ++++++++++++------ .../unit/extras/tasks/testdata/script.sh | 4 +- 3 files changed, 119 insertions(+), 87 deletions(-) diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py index 6e8dbcc21b..63bae594ed 100644 --- a/flytekit/extras/tasks/shell.py +++ b/flytekit/extras/tasks/shell.py @@ -1,7 +1,8 @@ +import collections import datetime import logging import os -import re +import string import subprocess import typing from dataclasses import dataclass @@ -30,46 +31,6 @@ class OutputLocation: location: typing.Union[os.PathLike, str] -def _stringify(v: typing.Any) -> str: - """ - Special cased return for the given value. Given the type returns the string version for the type. - Handles FlyteFile and FlyteDirectory specially. Downloads and returns the downloaded filepath - """ - if isinstance(v, FlyteFile): - v.download() - return v.path - if isinstance(v, FlyteDirectory): - v.download() - return v.path - if isinstance(v, datetime.datetime): - return v.isoformat() - return str(v) - - -def _interpolate(tmpl: str, regex: re.Pattern, validate_all_match: bool = True, **kwargs) -> str: - """ - Substitutes all templates that match the supplied regex - with the given inputs and returns the substituted string. The result is non destructive towards the given string. - """ - modified = tmpl - matched = set() - for match in regex.finditer(tmpl): - expr = match.groups()[0] - var = match.groups()[1] - if var not in kwargs: - raise ValueError(f"Variable {var} in Query (part of {expr}) not found in inputs {kwargs.keys()}") - matched.add(var) - val = kwargs[var] - # str conversion should be deliberate, with right conversion for each type - modified = modified.replace(expr, _stringify(val)) - - if validate_all_match: - if len(matched) < len(kwargs.keys()): - diff = set(kwargs.keys()).difference(matched) - raise ValueError(f"Extra Inputs have no matches in script template - missing {diff}") - return modified - - def _dummy_task_func(): """ A Fake function to satisfy the inner PythonTask requirements @@ -80,12 +41,51 @@ def _dummy_task_func(): T = typing.TypeVar("T") +class _PythonFStringInterpolizer: + """A class for interpolating scripts that use python string.format syntax""" + + class _Formatter(string.Formatter): + def format_field(self, value, format_spec): + """ + Special cased return for the given value. Given the type returns the string version for + the type. Handles FlyteFile and FlyteDirectory specially. + Downloads and returns the downloaded filepath. + """ + if isinstance(value, FlyteFile): + value.download() + return value.path + if isinstance(value, FlyteDirectory): + value.download() + return value.path + if isinstance(value, datetime.datetime): + return value.isoformat() + return super().format_field(value, format_spec) + + def interpolate( + self, + tmpl: str, + inputs: typing.Optional[typing.Dict[str, str]] = None, + outputs: typing.Optional[typing.Dict[str, str]] = None, + ) -> str: + """ + Interpolate python formatted string templates with variables from the input and output + argument dicts. The result is non destructive towards the given template string. + """ + inputs = inputs or {} + outputs = outputs or {} + reused_vars = inputs.keys() & outputs.keys() + if reused_vars: + raise ValueError(f"Variables {reused_vars} in Query cannot be shared between inputs and outputs.") + consolidated_args = collections.ChainMap(inputs, outputs) + try: + return self._Formatter().format(tmpl, **consolidated_args) + except KeyError as e: + raise ValueError(f"Variable {e} in Query not found in inputs {consolidated_args.keys()}") + + class ShellTask(PythonInstanceTask[T]): """ """ - _INPUT_REGEX = re.compile(r"({{\s*.inputs.(\w+)\s*}})", re.IGNORECASE) - _OUTPUT_REGEX = re.compile(r"({{\s*.outputs.(\w+)\s*}})", re.IGNORECASE) - def __init__( self, name: str, @@ -136,6 +136,7 @@ def __init__( self._script_file = script_file self._debug = debug self._output_locs = output_locs if output_locs else [] + self._interpolizer = _PythonFStringInterpolizer() outputs = self._validate_output_locs() super().__init__( name, @@ -184,12 +185,9 @@ def execute(self, **kwargs) -> typing.Any: outputs: typing.Dict[str, str] = {} if self._output_locs: for v in self._output_locs: - outputs[v.var] = _interpolate(v.location, self._INPUT_REGEX, validate_all_match=False, **kwargs) + outputs[v.var] = self._interpolizer.interpolate(v.location, inputs=kwargs) - gen_script = _interpolate(self._script, self._INPUT_REGEX, **kwargs) - # For outputs it is not necessary that all outputs are used in the script, some are implicit outputs - # for example gcc main.c will generate a.out automatically - gen_script = _interpolate(gen_script, self._OUTPUT_REGEX, validate_all_match=False, **outputs) + gen_script = self._interpolizer.interpolate(self._script, inputs=kwargs, outputs=outputs) if self._debug: print("\n==============================================\n") print(gen_script) diff --git a/tests/flytekit/unit/extras/tasks/test_shell.py b/tests/flytekit/unit/extras/tasks/test_shell.py index 39f114a9c7..2b76f7ad7f 100644 --- a/tests/flytekit/unit/extras/tasks/test_shell.py +++ b/tests/flytekit/unit/extras/tasks/test_shell.py @@ -1,9 +1,11 @@ import datetime import os import tempfile +from dataclasses import dataclass from subprocess import CalledProcessError import pytest +from dataclasses_json import dataclass_json from flytekit import kwtypes from flytekit.extras.tasks.shell import OutputLocation, ShellTask @@ -43,10 +45,10 @@ def test_input_substitution_primitive(): t = ShellTask( name="test", script=""" - set -ex - cat {{ .inputs.f }} - echo "Hello World {{ .inputs.y }} on {{ .inputs.j }}" - """, + set -ex + cat {f} + echo "Hello World {y} on {j}" + """, inputs=kwtypes(f=str, y=int, j=datetime.datetime), ) @@ -60,9 +62,9 @@ def test_input_substitution_files(): t = ShellTask( name="test", script=""" - cat {{ .inputs.f }} - echo "Hello World {{ .inputs.y }} on {{ .inputs.j }}" - """, + cat {f} + echo "Hello World {y} on {j}" + """, inputs=kwtypes(f=CSVFile, y=FlyteDirectory, j=datetime.datetime), ) @@ -70,20 +72,18 @@ def test_input_substitution_files(): def test_input_output_substitution_files(): - s = """ - cat {{ .inputs.f }} > {{ .outputs.y }} - """ + script = "cat {f} > {y}" t = ShellTask( name="test", debug=True, - script=s, + script=script, inputs=kwtypes(f=CSVFile), output_locs=[ - OutputLocation(var="y", var_type=FlyteFile, location="{{ .inputs.f }}.mod"), + OutputLocation(var="y", var_type=FlyteFile, location="{f}.mod"), ], ) - assert t.script == s + assert t.script == script contents = "1,2,3,4\n" with tempfile.TemporaryDirectory() as tmp: @@ -100,61 +100,95 @@ def test_input_output_substitution_files(): def test_input_single_output_substitution_files(): - s = """ - cat {{ .inputs.f }} >> {{ .outputs.y }} - echo "Hello World {{ .inputs.y }} on {{ .inputs.j }}" - """ + script = """ + cat {f} >> {z} + echo "Hello World {y} on {j}" + """ t = ShellTask( name="test", debug=True, - script=s, + script=script, inputs=kwtypes(f=CSVFile, y=FlyteDirectory, j=datetime.datetime), - output_locs=[OutputLocation(var="y", var_type=FlyteFile, location="{{ .inputs.f }}.pyc")], + output_locs=[OutputLocation(var="z", var_type=FlyteFile, location="{f}.pyc")], ) - assert t.script == s + assert t.script == script y = t(f=test_csv, y=testdata, j=datetime.datetime(2021, 11, 10, 12, 15, 0)) assert y.path[-4:] == ".pyc" -def test_input_output_extra_var_in_template(): +@pytest.mark.parametrize( + "script", + [ + ( + """ + cat {missing} >> {z} + echo "Hello World {y} on {j} - output {x}" + """ + ), + ( + """ + cat {f} {missing} >> {z} + echo "Hello World {y} on {j} - output {x}" + """ + ), + ], +) +def test_input_output_extra_and_missing_variables(script): t = ShellTask( name="test", debug=True, - script=""" - cat {{ .inputs.f }} {{ .inputs.missing }} >> {{ .outputs.y }} - echo "Hello World {{ .inputs.y }} on {{ .inputs.j }} - output {{.outputs.x}}" - """, + script=script, inputs=kwtypes(f=CSVFile, y=FlyteDirectory, j=datetime.datetime), output_locs=[ - OutputLocation(var="x", var_type=FlyteDirectory, location="{{ .inputs.y }}"), - OutputLocation(var="y", var_type=FlyteFile, location="{{ .inputs.f }}.pyc"), + OutputLocation(var="x", var_type=FlyteDirectory, location="{y}"), + OutputLocation(var="z", var_type=FlyteFile, location="{f}.pyc"), ], ) - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="missing"): t(f=test_csv, y=testdata, j=datetime.datetime(2021, 11, 10, 12, 15, 0)) -def test_input_output_extra_input(): +def test_cannot_reuse_variables_for_both_inputs_and_outputs(): t = ShellTask( name="test", debug=True, script=""" - cat {{ .inputs.missing }} >> {{ .outputs.y }} - echo "Hello World {{ .inputs.y }} on {{ .inputs.j }} - output {{.outputs.x}}" + cat {f} >> {y} + echo "Hello World {y} on {j}" """, inputs=kwtypes(f=CSVFile, y=FlyteDirectory, j=datetime.datetime), output_locs=[ - OutputLocation(var="x", var_type=FlyteDirectory, location="{{ .inputs.y }}"), - OutputLocation(var="y", var_type=FlyteFile, location="{{ .inputs.f }}.pyc"), + OutputLocation(var="y", var_type=FlyteFile, location="{f}.pyc"), ], ) - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="Variables {'y'} in Query"): t(f=test_csv, y=testdata, j=datetime.datetime(2021, 11, 10, 12, 15, 0)) +def test_can_use_complex_types_for_inputs_to_f_string_template(): + @dataclass_json + @dataclass + class InputArgs: + in_file: CSVFile + + t = ShellTask( + name="test", + debug=True, + script="""cat {input_args.in_file} >> {input_args.in_file}.tmp""", + inputs=kwtypes(input_args=InputArgs), + output_locs=[ + OutputLocation(var="x", var_type=FlyteFile, location="{input_args.in_file}.tmp"), + ], + ) + + input_args = InputArgs(FlyteFile(path=test_csv)) + x = t(input_args=input_args) + assert x.path[-4:] == ".tmp" + + def test_shell_script(): t = ShellTask( name="test2", @@ -162,8 +196,8 @@ def test_shell_script(): script_file=script_sh, inputs=kwtypes(f=CSVFile, y=FlyteDirectory, j=datetime.datetime), output_locs=[ - OutputLocation(var="x", var_type=FlyteDirectory, location="{{ .inputs.y }}"), - OutputLocation(var="y", var_type=FlyteFile, location="{{ .inputs.f }}.pyc"), + OutputLocation(var="x", var_type=FlyteDirectory, location="{y}"), + OutputLocation(var="z", var_type=FlyteFile, location="{f}.pyc"), ], ) diff --git a/tests/flytekit/unit/extras/tasks/testdata/script.sh b/tests/flytekit/unit/extras/tasks/testdata/script.sh index 22012ec3ae..1deb4c474a 100644 --- a/tests/flytekit/unit/extras/tasks/testdata/script.sh +++ b/tests/flytekit/unit/extras/tasks/testdata/script.sh @@ -2,5 +2,5 @@ set -ex -cat "{{ .inputs.f }}" >> "{{ .outputs.y }}" -echo "Hello World {{ .inputs.y }} on {{ .inputs.j }} - output {{.outputs.x}}" +cat "{f}" >> "{z}" +echo "Hello World {y} on {j} - output {x}" From de3aba30bcbc3dd8320934253cbbb66fc96a3900 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Tue, 18 Jan 2022 23:04:52 +0800 Subject: [PATCH 059/128] Expose configured RawOutputPrefix during execution (#813) * Expose configured RawOutputPrefix during execution Signed-off-by: Kevin Su * Remove sdk_runnable.py and spark_task.py Signed-off-by: Kevin Su Signed-off-by: maximsmol --- flytekit/bin/entrypoint.py | 1 + flytekit/core/context_manager.py | 14 ++++++++++++-- tests/flytekit/unit/core/test_type_hints.py | 1 + 3 files changed, 14 insertions(+), 2 deletions(-) diff --git a/flytekit/bin/entrypoint.py b/flytekit/bin/entrypoint.py index 754eab666d..7bec83346b 100644 --- a/flytekit/bin/entrypoint.py +++ b/flytekit/bin/entrypoint.py @@ -201,6 +201,7 @@ def setup_execution( ), logging=python_logging, tmp_dir=user_workspace_dir, + raw_output_prefix=ctx.file_access._raw_output_prefix, ) # TODO: Remove this check for flytekit 1.0 diff --git a/flytekit/core/context_manager.py b/flytekit/core/context_manager.py index 1db5e11d5e..a3ee0c2972 100644 --- a/flytekit/core/context_manager.py +++ b/flytekit/core/context_manager.py @@ -159,6 +159,7 @@ class Builder(object): execution_id: str attrs: typing.Dict[str, typing.Any] working_dir: typing.Union[os.PathLike, utils.AutoDeletingTempDir] + raw_output_prefix: str def __init__(self, current: typing.Optional[ExecutionParameters] = None): self.stats = current.stats if current else None @@ -167,6 +168,7 @@ def __init__(self, current: typing.Optional[ExecutionParameters] = None): self.execution_id = current.execution_id if current else None self.logging = current.logging if current else None self.attrs = current._attrs if current else {} + self.raw_output_prefix = current.raw_output_prefix if current else None def add_attr(self, key: str, v: typing.Any) -> ExecutionParameters.Builder: self.attrs[key] = v @@ -181,6 +183,7 @@ def build(self) -> ExecutionParameters: tmp_dir=self.working_dir, execution_id=self.execution_id, logging=self.logging, + raw_output_prefix=self.raw_output_prefix, **self.attrs, ) @@ -191,7 +194,7 @@ def new_builder(current: ExecutionParameters = None) -> Builder: def builder(self) -> Builder: return ExecutionParameters.Builder(current=self) - def __init__(self, execution_date, tmp_dir, stats, execution_id, logging, **kwargs): + def __init__(self, execution_date, tmp_dir, stats, execution_id, logging, raw_output_prefix, **kwargs): """ Args: execution_date: Date when the execution is running @@ -205,6 +208,7 @@ def __init__(self, execution_date, tmp_dir, stats, execution_id, logging, **kwar self._working_directory = tmp_dir self._execution_id = execution_id self._logging = logging + self._raw_output_prefix = raw_output_prefix # AutoDeletingTempDir's should be used with a with block, which creates upon entry self._attrs = kwargs # It is safe to recreate the Secrets Manager @@ -226,6 +230,10 @@ def logging(self) -> _logging: """ return self._logging + @property + def raw_output_prefix(self) -> str: + return self._raw_output_prefix + @property def working_directory(self) -> utils.AutoDeletingTempDir: """ @@ -895,14 +903,16 @@ def initialize(): # Note we use the SdkWorkflowExecution object purely for formatting into the ex:project:domain:name format users # are already acquainted with + default_context = FlyteContext(file_access=default_local_file_access_provider) default_user_space_params = ExecutionParameters( execution_id=str(WorkflowExecutionIdentifier.promote_from_model(default_execution_id)), execution_date=_datetime.datetime.utcnow(), stats=mock_stats.MockStats(), logging=_logging, tmp_dir=user_space_path, + raw_output_prefix=default_context.file_access._raw_output_prefix, ) - default_context = FlyteContext(file_access=default_local_file_access_provider) + default_context = default_context.with_execution_state( default_context.new_execution_state().with_params(user_space_params=default_user_space_params) ).build() diff --git a/tests/flytekit/unit/core/test_type_hints.py b/tests/flytekit/unit/core/test_type_hints.py index 367e71b467..37deb6d0d7 100644 --- a/tests/flytekit/unit/core/test_type_hints.py +++ b/tests/flytekit/unit/core/test_type_hints.py @@ -57,6 +57,7 @@ def test_default_wf_params_works(): def my_task(a: int): wf_params = flytekit.current_context() assert wf_params.execution_id == "ex:local:local:local" + assert "/tmp/flyte/" in wf_params.raw_output_prefix my_task(a=3) assert context_manager.FlyteContextManager.size() == 1 From 83fbe5043e9c200c1056594cc1f89afe06fc431f Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Tue, 18 Jan 2022 13:48:45 -0800 Subject: [PATCH 060/128] Add SecretsManager back to old import location (#820) Signed-off-by: Yee Hing Tong Signed-off-by: maximsmol --- flytekit/testing/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/flytekit/testing/__init__.py b/flytekit/testing/__init__.py index bb75358198..06b69612e5 100644 --- a/flytekit/testing/__init__.py +++ b/flytekit/testing/__init__.py @@ -16,4 +16,5 @@ """ +from flytekit.core.context_manager import SecretsManager from flytekit.core.testing import patch, task_mock From 2e602664983acd80ed0c38d49924e6d6687375bb Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Wed, 19 Jan 2022 10:40:42 -0800 Subject: [PATCH 061/128] Add some tests (#819) Signed-off-by: Yee Hing Tong Signed-off-by: maximsmol --- tests/flytekit/unit/core/test_imperative.py | 23 ++++++++ tests/flytekit/unit/core/test_shim_task.py | 63 +++++++++++++++++++++ 2 files changed, 86 insertions(+) create mode 100644 tests/flytekit/unit/core/test_shim_task.py diff --git a/tests/flytekit/unit/core/test_imperative.py b/tests/flytekit/unit/core/test_imperative.py index f120398306..2d34466862 100644 --- a/tests/flytekit/unit/core/test_imperative.py +++ b/tests/flytekit/unit/core/test_imperative.py @@ -212,6 +212,29 @@ def t2(a: typing.List[int]) -> int: assert wb() == [3, 6] +def test_imperative_tuples(): + @task + def t1() -> (int, str): + return 3, "three" + + @task + def t3(a: int, b: str) -> typing.Tuple[int, str]: + return a + 2, "world" + b + + wb = ImperativeWorkflow(name="my.workflow.a") + t1_node = wb.add_entity(t1) + t3_node = wb.add_entity(t3, a=t1_node.outputs["o0"], b=t1_node.outputs["o1"]) + wb.add_workflow_output("wf0", t3_node.outputs["o0"], python_type=int) + wb.add_workflow_output("wf1", t3_node.outputs["o1"], python_type=str) + res = wb() + assert res == (5, "worldthree") + + with pytest.raises(KeyError): + wb = ImperativeWorkflow(name="my.workflow.b") + t1_node = wb.add_entity(t1) + wb.add_entity(t3, a=t1_node.outputs["bad"], b=t1_node.outputs["o2"]) + + def test_call_normal(): @task def t1(a: int) -> (int, str): diff --git a/tests/flytekit/unit/core/test_shim_task.py b/tests/flytekit/unit/core/test_shim_task.py new file mode 100644 index 0000000000..b22f2e7349 --- /dev/null +++ b/tests/flytekit/unit/core/test_shim_task.py @@ -0,0 +1,63 @@ +import tempfile +from collections import OrderedDict + +import mock + +from flytekit import ContainerTask, kwtypes +from flytekit.core import context_manager +from flytekit.core.context_manager import Image, ImageConfig +from flytekit.core.python_customized_container_task import PythonCustomizedContainerTask, TaskTemplateResolver +from flytekit.core.utils import write_proto_to_file +from flytekit.tools.translator import get_serializable + +default_img = Image(name="default", fqn="test", tag="tag") +serialization_settings = context_manager.SerializationSettings( + project="project", + domain="domain", + version="version", + env=None, + image_config=ImageConfig(default_image=default_img, images=[default_img]), +) + + +class Placeholder(object): + ... + + +def test_resolver_load_task(): + # any task is fine, just copied one + square = ContainerTask( + name="square", + input_data_dir="/var/inputs", + output_data_dir="/var/outputs", + inputs=kwtypes(val=int), + outputs=kwtypes(out=int), + image="alpine", + command=["sh", "-c", "echo $(( {{.Inputs.val}} * {{.Inputs.val}} )) | tee /var/outputs/out"], + ) + + resolver = TaskTemplateResolver() + ts = get_serializable(OrderedDict(), serialization_settings, square) + with tempfile.NamedTemporaryFile() as f: + write_proto_to_file(ts.template.to_flyte_idl(), f.name) + # load_task should create an instance of the path to the object given, doesn't need to be a real executor + shim_task = resolver.load_task([f.name, f"{Placeholder.__module__}.Placeholder"]) + assert isinstance(shim_task.executor, Placeholder) + assert shim_task.task_template.id.name == "square" + assert shim_task.task_template.interface.inputs["val"] is not None + assert shim_task.task_template.interface.outputs["out"] is not None + + +@mock.patch("flytekit.core.python_customized_container_task.PythonCustomizedContainerTask.get_config") +@mock.patch("flytekit.core.python_customized_container_task.PythonCustomizedContainerTask.get_custom") +def test_serialize_to_model(mock_custom, mock_config): + mock_custom.return_value = {"a": "custom"} + mock_config.return_value = {"a": "config"} + ct = PythonCustomizedContainerTask( + name="mytest", task_config=None, container_image="someimage", executor_type=Placeholder + ) + tt = ct.serialize_to_model(serialization_settings) + assert tt.container.image == "someimage" + assert len(tt.config) == 1 + assert tt.id.name == "mytest" + assert len(tt.custom) == 1 From 057c130164c72da431d9b374bbe3e93de476b18f Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Fri, 21 Jan 2022 07:25:29 +0800 Subject: [PATCH 062/128] Fixed flaky spark test (#821) Signed-off-by: Kevin Su Signed-off-by: maximsmol --- plugins/flytekit-spark/flytekitplugins/spark/task.py | 1 + 1 file changed, 1 insertion(+) diff --git a/plugins/flytekit-spark/flytekitplugins/spark/task.py b/plugins/flytekit-spark/flytekitplugins/spark/task.py index 37ae03e913..c05dbc2ea0 100644 --- a/plugins/flytekit-spark/flytekitplugins/spark/task.py +++ b/plugins/flytekit-spark/flytekitplugins/spark/task.py @@ -110,6 +110,7 @@ def pre_execute(self, user_params: ExecutionParameters) -> ExecutionParameters: # If either of above cases is not true, then we are in local execution of this task # Add system spark-conf for local/notebook based execution. spark_conf = _pyspark.SparkConf() + spark_conf.set("spark.driver.bindAddress", "127.0.0.1") for k, v in self.task_config.spark_conf.items(): spark_conf.set(k, v) # In local execution, propagate PYTHONPATH to executors too. This makes the spark From 3248ddf5efdf3d1db5e9e290217d8a48a1ce54ba Mon Sep 17 00:00:00 2001 From: Snyk bot Date: Fri, 21 Jan 2022 18:16:58 +0000 Subject: [PATCH 063/128] fix: plugins/flytekit-greatexpectations/requirements.txt to reduce vulnerabilities (#823) The following vulnerabilities are fixed by pinning transitive dependencies: - https://snyk.io/vuln/SNYK-PYTHON-IPYTHON-2348630 Signed-off-by: maximsmol --- plugins/flytekit-greatexpectations/requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/plugins/flytekit-greatexpectations/requirements.txt b/plugins/flytekit-greatexpectations/requirements.txt index f5aa2f1471..5bdbc66553 100644 --- a/plugins/flytekit-greatexpectations/requirements.txt +++ b/plugins/flytekit-greatexpectations/requirements.txt @@ -88,7 +88,7 @@ ipykernel==6.5.0 # via # ipywidgets # notebook -ipython==7.29.0 +ipython==7.31.1 # via # ipykernel # ipywidgets From f2a0405cfc9893639f74568c7ca87f408d3192ae Mon Sep 17 00:00:00 2001 From: Snyk bot Date: Fri, 21 Jan 2022 18:17:22 +0000 Subject: [PATCH 064/128] fix: plugins/flytekit-papermill/requirements.txt to reduce vulnerabilities (#825) The following vulnerabilities are fixed by pinning transitive dependencies: - https://snyk.io/vuln/SNYK-PYTHON-IPYTHON-2348630 Signed-off-by: maximsmol --- plugins/flytekit-papermill/requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/plugins/flytekit-papermill/requirements.txt b/plugins/flytekit-papermill/requirements.txt index 9c5a82d050..a4b8b61e3d 100644 --- a/plugins/flytekit-papermill/requirements.txt +++ b/plugins/flytekit-papermill/requirements.txt @@ -79,7 +79,7 @@ importlib-metadata==4.8.2 # via keyring ipykernel==6.5.0 # via flytekitplugins-papermill -ipython==7.29.0 +ipython==7.31.1 # via ipykernel ipython-genutils==0.2.0 # via nbformat From 85b5057c4e519684338bcc8c2f0803fbe667a1a1 Mon Sep 17 00:00:00 2001 From: Snyk bot Date: Fri, 21 Jan 2022 18:17:42 +0000 Subject: [PATCH 065/128] fix: requirements-spark2.txt to reduce vulnerabilities (#826) The following vulnerabilities are fixed by pinning transitive dependencies: - https://snyk.io/vuln/SNYK-PYTHON-IPYTHON-2348630 Signed-off-by: maximsmol --- requirements-spark2.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-spark2.txt b/requirements-spark2.txt index ecc11d57ea..bdef3764f0 100644 --- a/requirements-spark2.txt +++ b/requirements-spark2.txt @@ -101,7 +101,7 @@ inotify_simple==1.2.1 # via sagemaker-training ipykernel==5.5.6 # via flytekit -ipython==7.30.1 +ipython==7.31.1 # via ipykernel ipython-genutils==0.2.0 # via From 9167cffe7c3c8d8b8127e8f79b7bda333fc2b8f7 Mon Sep 17 00:00:00 2001 From: Snyk bot Date: Fri, 21 Jan 2022 18:18:06 +0000 Subject: [PATCH 066/128] fix: requirements.txt to reduce vulnerabilities (#824) The following vulnerabilities are fixed by pinning transitive dependencies: - https://snyk.io/vuln/SNYK-PYTHON-IPYTHON-2348630 Signed-off-by: maximsmol --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index 9ab1213845..38f885cec8 100644 --- a/requirements.txt +++ b/requirements.txt @@ -99,7 +99,7 @@ inotify_simple==1.2.1 # via sagemaker-training ipykernel==5.5.6 # via flytekit -ipython==7.30.1 +ipython==7.31.1 # via ipykernel ipython-genutils==0.2.0 # via From 193500660ae99dcdf737395d595ba407e11d85fa Mon Sep 17 00:00:00 2001 From: Ketan Umare <16888709+kumare3@users.noreply.github.com> Date: Fri, 21 Jan 2022 15:32:33 -0800 Subject: [PATCH 067/128] Intratask checkpointing (#771) * wip - intratask checkpointing Signed-off-by: Ketan Umare * sync checkpointer with tests Signed-off-by: Ketan Umare * Checkpinter in entrypoint Signed-off-by: Ketan Umare * checkpoint in progress Signed-off-by: Ketan Umare * wip Signed-off-by: Ketan Umare * Intratask checkpointer Signed-off-by: Ketan Umare * Checkpoint updated Signed-off-by: Ketan Umare * Intra-task checkpointing Signed-off-by: Ketan Umare * Test and entrypoint updated Signed-off-by: Ketan Umare * lint fixed Signed-off-by: Ketan Umare * test fixes Signed-off-by: Ketan Umare * fmt Signed-off-by: Ketan Umare * updated entrypoint Signed-off-by: Ketan Umare * updated Signed-off-by: Ketan Umare * update Signed-off-by: Ketan Umare * print Signed-off-by: Ketan Umare * updated Signed-off-by: Ketan Umare * SyncCheckpointer working Signed-off-by: Ketan Umare * updated Signed-off-by: Ketan Umare * update Signed-off-by: Ketan Umare * fixed import problems Signed-off-by: Ketan Umare * fixed test Signed-off-by: Ketan Umare * fixed imports Signed-off-by: Ketan Umare * fixed lints and errors Signed-off-by: Ketan Umare * lint fix Signed-off-by: Ketan Umare * addressed comments Signed-off-by: Ketan Umare Signed-off-by: maximsmol --- flytekit/bin/entrypoint.py | 122 ++++++++++---- flytekit/core/base_task.py | 3 + flytekit/core/checkpointer.py | 157 ++++++++++++++++++ flytekit/core/context_manager.py | 31 +++- flytekit/core/map_task.py | 4 + flytekit/core/python_auto_container.py | 7 +- flytekit/core/utils.py | 26 +-- plugins/flytekit-k8s-pod/tests/test_pod.py | 16 ++ .../unit/bin/test_python_entrypoint.py | 17 +- tests/flytekit/unit/core/test_checkpoint.py | 117 +++++++++++++ tests/flytekit/unit/core/test_checkpointer.py | 67 ++++++++ tests/flytekit/unit/core/test_map_task.py | 4 + .../unit/core/test_python_auto_container.py | 45 ++++- 13 files changed, 561 insertions(+), 55 deletions(-) create mode 100644 flytekit/core/checkpointer.py create mode 100644 tests/flytekit/unit/core/test_checkpoint.py create mode 100644 tests/flytekit/unit/core/test_checkpointer.py diff --git a/flytekit/bin/entrypoint.py b/flytekit/bin/entrypoint.py index 7bec83346b..4fc9f69d3e 100644 --- a/flytekit/bin/entrypoint.py +++ b/flytekit/bin/entrypoint.py @@ -4,7 +4,7 @@ import os as _os import pathlib import traceback as _traceback -from typing import List +from typing import List, Optional import click as _click from flyteidl.core import literals_pb2 as _literals_pb2 @@ -16,6 +16,7 @@ from flytekit.core import constants as _constants from flytekit.core import utils from flytekit.core.base_task import IgnoreOutputs, PythonTask +from flytekit.core.checkpointer import SyncCheckpoint from flytekit.core.context_manager import ( ExecutionParameters, ExecutionState, @@ -164,8 +165,10 @@ def _dispatch_execute( @contextlib.contextmanager def setup_execution( raw_output_data_prefix: str, - dynamic_addl_distro: str = None, - dynamic_dest_dir: str = None, + checkpoint_path: Optional[str] = None, + prev_checkpoint: Optional[str] = None, + dynamic_addl_distro: Optional[str] = None, + dynamic_dest_dir: Optional[str] = None, ): ctx = FlyteContextManager.current_context() @@ -175,6 +178,11 @@ def setup_execution( pathlib.Path(user_workspace_dir).mkdir(parents=True, exist_ok=True) from flytekit import __version__ as _api_version + checkpointer = None + if checkpoint_path is not None: + checkpointer = SyncCheckpoint(checkpoint_dest=checkpoint_path, checkpoint_src=prev_checkpoint) + logger.debug(f"Checkpointer created with source {prev_checkpoint} and dest {checkpoint_path}") + execution_parameters = ExecutionParameters( execution_id=_identifier.WorkflowExecutionIdentifier( project=_internal_config.EXECUTION_PROJECT.get(), @@ -202,6 +210,7 @@ def setup_execution( logging=python_logging, tmp_dir=user_workspace_dir, raw_output_prefix=ctx.file_access._raw_output_prefix, + checkpoint=checkpointer, ) # TODO: Remove this check for flytekit 1.0 @@ -266,14 +275,16 @@ def _handle_annotated_task( @_scopes.system_entry_point def _execute_task( - inputs, - output_prefix, - raw_output_data_prefix, - test, + inputs: str, + output_prefix: str, + test: bool, + raw_output_data_prefix: str, resolver: str, resolver_args: List[str], - dynamic_addl_distro: str = None, - dynamic_dest_dir: str = None, + checkpoint_path: Optional[str] = None, + prev_checkpoint: Optional[str] = None, + dynamic_addl_distro: Optional[str] = None, + dynamic_dest_dir: Optional[str] = None, ): """ This function should be called for new API tasks (those only available in 0.16 and later that leverage Python @@ -302,7 +313,13 @@ def _execute_task( raise Exception("cannot be <1") with _TemporaryConfiguration(_internal_config.CONFIGURATION_PATH.get()): - with setup_execution(raw_output_data_prefix, dynamic_addl_distro, dynamic_dest_dir) as ctx: + with setup_execution( + raw_output_data_prefix, + checkpoint_path=checkpoint_path, + prev_checkpoint=prev_checkpoint, + dynamic_addl_distro=dynamic_addl_distro, + dynamic_dest_dir=dynamic_dest_dir, + ) as ctx: resolver_obj = load_object_from_module(resolver) # Use the resolver to load the actual task object _task_def = resolver_obj.load_task(loader_args=resolver_args) @@ -321,16 +338,20 @@ def _execute_map_task( raw_output_data_prefix, max_concurrency, test, - dynamic_addl_distro: str, - dynamic_dest_dir: str, resolver: str, resolver_args: List[str], + checkpoint_path: Optional[str] = None, + prev_checkpoint: Optional[str] = None, + dynamic_addl_distro: Optional[str] = None, + dynamic_dest_dir: Optional[str] = None, ): if len(resolver_args) < 1: raise Exception(f"Resolver args cannot be <1, got {resolver_args}") with _TemporaryConfiguration(_internal_config.CONFIGURATION_PATH.get()): - with setup_execution(raw_output_data_prefix, dynamic_addl_distro, dynamic_dest_dir) as ctx: + with setup_execution( + raw_output_data_prefix, checkpoint_path, prev_checkpoint, dynamic_addl_distro, dynamic_dest_dir + ) as ctx: resolver_obj = load_object_from_module(resolver) # Use the resolver to load the actual task object _task_def = resolver_obj.load_task(loader_args=resolver_args) @@ -352,6 +373,22 @@ def _execute_map_task( _handle_annotated_task(ctx, map_task, inputs, output_prefix) +def normalize_inputs( + raw_output_data_prefix: Optional[str], checkpoint_path: Optional[str], prev_checkpoint: Optional[str] +): + # Backwards compatibility - if Propeller hasn't filled this in, then it'll come through here as the original + # template string, so let's explicitly set it to None so that the downstream functions will know to fall back + # to the original shard formatter/prefix config. + if raw_output_data_prefix == "{{.rawOutputDataPrefix}}": + raw_output_data_prefix = None + if checkpoint_path == "{{.checkpointOutputPrefix}}": + checkpoint_path = None + if prev_checkpoint == "{{.prevCheckpointPrefix}}" or prev_checkpoint == "" or prev_checkpoint == '""': + prev_checkpoint = None + + return raw_output_data_prefix, checkpoint_path, prev_checkpoint + + @_click.group() def _pass_through(): pass @@ -361,6 +398,8 @@ def _pass_through(): @_click.option("--inputs", required=True) @_click.option("--output-prefix", required=True) @_click.option("--raw-output-data-prefix", required=False) +@_click.option("--checkpoint-path", required=False) +@_click.option("--prev-checkpoint", required=False) @_click.option("--test", is_flag=True) @_click.option("--dynamic-addl-distro", required=False) @_click.option("--dynamic-dest-dir", required=False) @@ -375,6 +414,8 @@ def execute_task_cmd( output_prefix, raw_output_data_prefix, test, + prev_checkpoint, + checkpoint_path, dynamic_addl_distro, dynamic_dest_dir, resolver, @@ -383,26 +424,27 @@ def execute_task_cmd( logger.info(get_version_message()) # We get weird errors if there are no click echo messages at all, so emit an empty string so that unit tests pass. _click.echo("") - # Backwards compatibility - if Propeller hasn't filled this in, then it'll come through here as the original - # template string, so let's explicitly set it to None so that the downstream functions will know to fall back - # to the original shard formatter/prefix config. - if raw_output_data_prefix == "{{.rawOutputDataPrefix}}": - raw_output_data_prefix = None + raw_output_data_prefix, checkpoint_path, prev_checkpoint = normalize_inputs( + raw_output_data_prefix, checkpoint_path, prev_checkpoint + ) # For new API tasks (as of 0.16.x), we need to call a different function. # Use the presence of the resolver to differentiate between old API tasks and new API tasks # The addition of a new top-level command seemed out of scope at the time of this writing to pursue given how # pervasive this top level command already (plugins mostly). + logger.debug(f"Running task execution with resolver {resolver}...") _execute_task( - inputs, - output_prefix, - raw_output_data_prefix, - test, - resolver, - resolver_args, - dynamic_addl_distro, - dynamic_dest_dir, + inputs=inputs, + output_prefix=output_prefix, + raw_output_data_prefix=raw_output_data_prefix, + test=test, + resolver=resolver, + resolver_args=resolver_args, + dynamic_addl_distro=dynamic_addl_distro, + dynamic_dest_dir=dynamic_dest_dir, + checkpoint_path=checkpoint_path, + prev_checkpoint=prev_checkpoint, ) @@ -446,6 +488,8 @@ def fast_execute_task_cmd(additional_distribution, dest_dir, task_execute_cmd): @_click.option("--dynamic-addl-distro", required=False) @_click.option("--dynamic-dest-dir", required=False) @_click.option("--resolver", required=True) +@_click.option("--checkpoint-path", required=False) +@_click.option("--prev-checkpoint", required=False) @_click.argument( "resolver-args", type=_click.UNPROCESSED, @@ -461,19 +505,27 @@ def map_execute_task_cmd( dynamic_dest_dir, resolver, resolver_args, + prev_checkpoint, + checkpoint_path, ): logger.info(get_version_message()) + raw_output_data_prefix, checkpoint_path, prev_checkpoint = normalize_inputs( + raw_output_data_prefix, checkpoint_path, prev_checkpoint + ) + _execute_map_task( - inputs, - output_prefix, - raw_output_data_prefix, - max_concurrency, - test, - dynamic_addl_distro, - dynamic_dest_dir, - resolver, - resolver_args, + inputs=inputs, + output_prefix=output_prefix, + raw_output_data_prefix=raw_output_data_prefix, + max_concurrency=max_concurrency, + test=test, + dynamic_addl_distro=dynamic_addl_distro, + dynamic_dest_dir=dynamic_dest_dir, + resolver=resolver, + resolver_args=resolver_args, + checkpoint_path=checkpoint_path, + prev_checkpoint=prev_checkpoint, ) diff --git a/flytekit/core/base_task.py b/flytekit/core/base_task.py index f590ba2033..c5c35373d8 100644 --- a/flytekit/core/base_task.py +++ b/flytekit/core/base_task.py @@ -263,6 +263,9 @@ def local_execute(self, ctx: FlyteContext, **kwargs) -> Union[Tuple[Promise], Pr else: logger.info("Cache hit") else: + es = ctx.execution_state + b = es.user_space_params.with_task_sandbox() + ctx = ctx.current_context().with_execution_state(es.with_params(user_space_params=b.build())).build() outputs_literal_map = self.dispatch_execute(ctx, input_literal_map) outputs_literals = outputs_literal_map.literals diff --git a/flytekit/core/checkpointer.py b/flytekit/core/checkpointer.py new file mode 100644 index 0000000000..bd17b7748b --- /dev/null +++ b/flytekit/core/checkpointer.py @@ -0,0 +1,157 @@ +import io +import tempfile +import typing +from abc import abstractmethod +from pathlib import Path + + +class Checkpoint(object): + """ + Base class for Checkpoint system. Checkpoint system allows reading and writing custom checkpoints from user + scripts + """ + + @abstractmethod + def prev_exists(self) -> bool: + raise NotImplementedError("Use one of the derived classes") + + @abstractmethod + def restore(self, path: typing.Union[Path, str]) -> typing.Optional[Path]: + """ + Given a path, if a previous checkpoint exists, will be downloaded to this path. + If download is successful the downloaded path is returned + + .. note: + + Download will not be performed, if the checkpoint was previously restored. The method will return the + previously downloaded path. + + """ + raise NotImplementedError("Use one of the derived classes") + + @abstractmethod + def save(self, cp: typing.Union[Path, str, io.BufferedReader]): + """ + Args: + cp: Checkpoint file (path, str path or a io.BufferedReader) + + Usage: If you have a io.BufferedReader then the following should work + + .. code-block: python + + with input_file.open(mode="rb") as b: + checkpointer.save(b) + """ + raise NotImplementedError("Use one of the derived classes") + + @abstractmethod + def read(self) -> typing.Optional[bytes]: + """ + This should only be used if there is a singular checkpoint file written. If more than one checkpoint file is + found, this will raise a ValueError + """ + raise NotImplementedError("Use one of the derived classes") + + @abstractmethod + def write(self, b: bytes): + """ + This will overwrite the checkpoint. It can be retrieved using read or restore + """ + raise NotImplementedError("Use one of the derived classes") + + +class SyncCheckpoint(Checkpoint): + """ + This class is NOT THREAD-SAFE! + Sync Checkpoint, will synchronously checkpoint a user given file or folder. + It will also synchronously download / restore previous checkpoints, when restore is invoked. + + TODO: Implement an async checkpoint system + """ + + SRC_LOCAL_FOLDER = "prev_cp" + TMP_DST_PATH = "_dst_cp" + + def __init__(self, checkpoint_dest: str, checkpoint_src: typing.Optional[str] = None): + """ + Args: + checkpoint_src: If a previous checkpoint should exist, this path should be set to the folder that contains the checkpoint information + checkpoint_dest: Location where the new checkpoint should be copied to + """ + self._checkpoint_dest = checkpoint_dest + self._checkpoint_src = checkpoint_src if checkpoint_src and checkpoint_src != "" else None + self._td = tempfile.TemporaryDirectory() + self._prev_download_path = None + + def __del__(self): + self._td.cleanup() + + def prev_exists(self) -> bool: + return self._checkpoint_src is not None + + def restore(self, path: typing.Optional[typing.Union[Path, str]] = None) -> typing.Optional[Path]: + + # We have to lazy load, until we fix the imports + from flytekit.core.context_manager import FlyteContextManager + + if self._checkpoint_src is None or self._checkpoint_src == "": + return None + + if self._prev_download_path: + return self._prev_download_path + + if path is None: + p = Path(self._td.name) + path = p.joinpath(self.SRC_LOCAL_FOLDER) + path.mkdir() + elif isinstance(path, str): + path = Path(path) + + if not path.is_dir(): + raise ValueError("Checkpoints can be restored to a directory only.") + + FlyteContextManager.current_context().file_access.download_directory(self._checkpoint_src, str(path)) + self._prev_download_path = path + return self._prev_download_path + + def save(self, cp: typing.Union[Path, str, io.BufferedReader]): + # We have to lazy load, until we fix the imports + from flytekit.core.context_manager import FlyteContextManager + + fa = FlyteContextManager.current_context().file_access + if isinstance(cp, (Path, str)): + if isinstance(cp, str): + cp = Path(cp) + if cp.is_dir(): + fa.upload_directory(str(cp), self._checkpoint_dest) + else: + fname = cp.stem + cp.suffix + rpath = fa._default_remote.construct_path(False, False, self._checkpoint_dest, fname) + fa.upload(str(cp), rpath) + return + + if not isinstance(cp, io.IOBase): + raise ValueError(f"Only a valid path or IOBase type (reader) should be provided, received {type(cp)}") + + p = Path(self._td.name) + dest_cp = p.joinpath(self.TMP_DST_PATH) + with dest_cp.open("wb") as f: + f.write(cp.read()) + + rpath = fa._default_remote.construct_path(False, False, self._checkpoint_dest, self.TMP_DST_PATH) + fa.upload(str(dest_cp), rpath) + + def read(self) -> typing.Optional[bytes]: + p = self.restore() + if p is None: + return None + files = list(p.iterdir()) + if len(files) == 0 or len(files) > 1: + raise ValueError(f"Expected exactly one checkpoint - found {len(files)}") + f = files[0] + return f.read_bytes() + + def write(self, b: bytes): + f = io.BytesIO(b) + f = typing.cast(io.BufferedReader, f) + self.save(f) diff --git a/flytekit/core/context_manager.py b/flytekit/core/context_manager.py index a3ee0c2972..70af22f1e1 100644 --- a/flytekit/core/context_manager.py +++ b/flytekit/core/context_manager.py @@ -19,6 +19,7 @@ import os import pathlib import re +import tempfile import traceback import typing from contextlib import contextmanager @@ -34,6 +35,7 @@ from flytekit.configuration import sdk as _sdk_config from flytekit.configuration import secrets from flytekit.core import mock_stats, utils +from flytekit.core.checkpointer import Checkpoint, SyncCheckpoint from flytekit.core.data_persistence import FileAccessProvider, default_local_file_access_provider from flytekit.core.node import Node from flytekit.interfaces.cli_identifiers import WorkflowExecutionIdentifier @@ -159,6 +161,7 @@ class Builder(object): execution_id: str attrs: typing.Dict[str, typing.Any] working_dir: typing.Union[os.PathLike, utils.AutoDeletingTempDir] + checkpoint: typing.Optional[Checkpoint] raw_output_prefix: str def __init__(self, current: typing.Optional[ExecutionParameters] = None): @@ -167,6 +170,7 @@ def __init__(self, current: typing.Optional[ExecutionParameters] = None): self.working_dir = current.working_directory if current else None self.execution_id = current.execution_id if current else None self.logging = current.logging if current else None + self.checkpoint = current._checkpoint if current else None self.attrs = current._attrs if current else {} self.raw_output_prefix = current.raw_output_prefix if current else None @@ -183,6 +187,7 @@ def build(self) -> ExecutionParameters: tmp_dir=self.working_dir, execution_id=self.execution_id, logging=self.logging, + checkpoint=self.checkpoint, raw_output_prefix=self.raw_output_prefix, **self.attrs, ) @@ -191,10 +196,26 @@ def build(self) -> ExecutionParameters: def new_builder(current: ExecutionParameters = None) -> Builder: return ExecutionParameters.Builder(current=current) + def with_task_sandbox(self) -> Builder: + prefix = self.working_directory + if isinstance(self.working_directory, utils.AutoDeletingTempDir): + prefix = self.working_directory.name + task_sandbox_dir = tempfile.mkdtemp(prefix=prefix) + p = pathlib.Path(task_sandbox_dir) + cp_dir = p.joinpath("__cp") + cp_dir.mkdir(exist_ok=True) + cp = SyncCheckpoint(checkpoint_dest=str(cp_dir)) + b = self.new_builder(self) + b.checkpoint = cp + b.working_dir = task_sandbox_dir + return b + def builder(self) -> Builder: return ExecutionParameters.Builder(current=self) - def __init__(self, execution_date, tmp_dir, stats, execution_id, logging, raw_output_prefix, **kwargs): + def __init__( + self, execution_date, tmp_dir, stats, execution_id, logging, raw_output_prefix, checkpoint=None, **kwargs + ): """ Args: execution_date: Date when the execution is running @@ -202,6 +223,7 @@ def __init__(self, execution_date, tmp_dir, stats, execution_id, logging, raw_ou stats: handle to emit stats execution_id: Identifier for the xecution logging: handle to logging + checkpoint: Checkpoint Handle to the configured checkpoint system """ self._stats = stats self._execution_date = execution_date @@ -213,6 +235,7 @@ def __init__(self, execution_date, tmp_dir, stats, execution_id, logging, raw_ou self._attrs = kwargs # It is safe to recreate the Secrets Manager self._secrets_manager = SecretsManager() + self._checkpoint = checkpoint @property def stats(self) -> taggable.TaggableStats: @@ -274,6 +297,12 @@ def execution_id(self) -> str: def secrets(self) -> SecretsManager: return self._secrets_manager + @property + def checkpoint(self) -> Checkpoint: + if self._checkpoint is None: + raise NotImplementedError("Checkpointing is not available, please check the version of the platform.") + return self._checkpoint + def __getattr__(self, attr_name: str) -> typing.Any: """ This houses certain task specific context. For example in Spark, it houses the SparkSession, etc diff --git a/flytekit/core/map_task.py b/flytekit/core/map_task.py index 29838cffcd..f760be5d3c 100644 --- a/flytekit/core/map_task.py +++ b/flytekit/core/map_task.py @@ -78,6 +78,10 @@ def get_command(self, settings: SerializationSettings) -> List[str]: "{{.outputPrefix}}", "--raw-output-data-prefix", "{{.rawOutputDataPrefix}}", + "--checkpoint-path", + "{{.checkpointOutputPrefix}}", + "--prev-checkpoint", + "{{.prevCheckpointPrefix}}", "--resolver", self._run_task.task_resolver.location, "--", diff --git a/flytekit/core/python_auto_container.py b/flytekit/core/python_auto_container.py index 0226760f08..c5f8413fea 100644 --- a/flytekit/core/python_auto_container.py +++ b/flytekit/core/python_auto_container.py @@ -2,6 +2,7 @@ import importlib import re +from abc import ABC from typing import Callable, Dict, List, Optional, TypeVar from flytekit.core.base_task import PythonTask, TaskResolverMixin @@ -17,7 +18,7 @@ T = TypeVar("T") -class PythonAutoContainerTask(PythonTask[T], metaclass=FlyteTrackedABC): +class PythonAutoContainerTask(PythonTask[T], ABC, metaclass=FlyteTrackedABC): """ A Python AutoContainer task should be used as the base for all extensions that want the user's code to be in the container and the container information to be automatically captured. @@ -119,6 +120,10 @@ def get_default_command(self, settings: SerializationSettings) -> List[str]: "{{.outputPrefix}}", "--raw-output-data-prefix", "{{.rawOutputDataPrefix}}", + "--checkpoint-path", + "{{.checkpointOutputPrefix}}", + "--prev-checkpoint", + "{{.prevCheckpointPrefix}}", "--resolver", self.task_resolver.location, "--", diff --git a/flytekit/core/utils.py b/flytekit/core/utils.py index 71dd45f581..db3d4b573b 100644 --- a/flytekit/core/utils.py +++ b/flytekit/core/utils.py @@ -5,7 +5,7 @@ import time as _time from hashlib import sha224 as _sha224 from pathlib import Path -from typing import Dict, List +from typing import Dict, List, Optional from flytekit.configuration import resources as _resource_config from flytekit.models import task as _task_models @@ -53,18 +53,18 @@ def _get_container_definition( image: str, command: List[str], args: List[str], - data_loading_config: _task_models.DataLoadingConfig, - storage_request: str = None, - ephemeral_storage_request: str = None, - cpu_request: str = None, - gpu_request: str = None, - memory_request: str = None, - storage_limit: str = None, - ephemeral_storage_limit: str = None, - cpu_limit: str = None, - gpu_limit: str = None, - memory_limit: str = None, - environment: Dict[str, str] = None, + data_loading_config: Optional[_task_models.DataLoadingConfig] = None, + storage_request: Optional[str] = None, + ephemeral_storage_request: Optional[str] = None, + cpu_request: Optional[str] = None, + gpu_request: Optional[str] = None, + memory_request: Optional[str] = None, + storage_limit: Optional[str] = None, + ephemeral_storage_limit: Optional[str] = None, + cpu_limit: Optional[str] = None, + gpu_limit: Optional[str] = None, + memory_limit: Optional[str] = None, + environment: Optional[Dict[str, str]] = None, ) -> _task_models.Container: storage_limit = storage_limit or _resource_config.DEFAULT_STORAGE_LIMIT.get() storage_request = storage_request or _resource_config.DEFAULT_STORAGE_REQUEST.get() diff --git a/plugins/flytekit-k8s-pod/tests/test_pod.py b/plugins/flytekit-k8s-pod/tests/test_pod.py index 82a0bcf6f8..9677f51211 100644 --- a/plugins/flytekit-k8s-pod/tests/test_pod.py +++ b/plugins/flytekit-k8s-pod/tests/test_pod.py @@ -69,6 +69,10 @@ def simple_pod_task(i: int): "{{.outputPrefix}}", "--raw-output-data-prefix", "{{.rawOutputDataPrefix}}", + "--checkpoint-path", + "{{.checkpointOutputPrefix}}", + "--prev-checkpoint", + "{{.prevCheckpointPrefix}}", "--resolver", "flytekit.core.python_auto_container.default_task_resolver", "--", @@ -134,6 +138,10 @@ def simple_pod_task(i: int): "{{.outputPrefix}}", "--raw-output-data-prefix", "{{.rawOutputDataPrefix}}", + "--checkpoint-path", + "{{.checkpointOutputPrefix}}", + "--prev-checkpoint", + "{{.prevCheckpointPrefix}}", "--resolver", "flytekit.core.python_auto_container.default_task_resolver", "--", @@ -321,6 +329,10 @@ def simple_pod_task(i: int): "{{.outputPrefix}}", "--raw-output-data-prefix", "{{.rawOutputDataPrefix}}", + "--checkpoint-path", + "{{.checkpointOutputPrefix}}", + "--prev-checkpoint", + "{{.prevCheckpointPrefix}}", "--resolver", "flytekit.core.python_auto_container.default_task_resolver", "--", @@ -367,6 +379,10 @@ def simple_pod_task(i: int): "{{.outputPrefix}}", "--raw-output-data-prefix", "{{.rawOutputDataPrefix}}", + "--checkpoint-path", + "{{.checkpointOutputPrefix}}", + "--prev-checkpoint", + "{{.prevCheckpointPrefix}}", "--resolver", "flytekit.core.python_auto_container.default_task_resolver", "--", diff --git a/tests/flytekit/unit/bin/test_python_entrypoint.py b/tests/flytekit/unit/bin/test_python_entrypoint.py index 3a71b567d8..4ccb3c8bcd 100644 --- a/tests/flytekit/unit/bin/test_python_entrypoint.py +++ b/tests/flytekit/unit/bin/test_python_entrypoint.py @@ -5,7 +5,7 @@ import pytest from flyteidl.core.errors_pb2 import ErrorDocument -from flytekit.bin.entrypoint import _dispatch_execute, setup_execution +from flytekit.bin.entrypoint import _dispatch_execute, normalize_inputs, setup_execution from flytekit.core import context_manager from flytekit.core.base_task import IgnoreOutputs from flytekit.core.dynamic_workflow_task import dynamic @@ -284,8 +284,19 @@ def test_setup_bad_prefix(): def test_setup_cloud_prefix(): - with setup_execution("s3://") as ctx: + with setup_execution("s3://", checkpoint_path=None, prev_checkpoint=None) as ctx: assert isinstance(ctx.file_access._default_remote, S3Persistence) - with setup_execution("gs://") as ctx: + with setup_execution("gs://", checkpoint_path=None, prev_checkpoint=None) as ctx: assert isinstance(ctx.file_access._default_remote, GCSPersistence) + + +def test_normalize_inputs(): + assert normalize_inputs("{{.rawOutputDataPrefix}}", "{{.checkpointOutputPrefix}}", "{{.prevCheckpointPrefix}}") == ( + None, + None, + None, + ) + assert normalize_inputs("/raw", "/cp1", '""') == ("/raw", "/cp1", None) + assert normalize_inputs("/raw", "/cp1", "") == ("/raw", "/cp1", None) + assert normalize_inputs("/raw", "/cp1", "/prev") == ("/raw", "/cp1", "/prev") diff --git a/tests/flytekit/unit/core/test_checkpoint.py b/tests/flytekit/unit/core/test_checkpoint.py new file mode 100644 index 0000000000..0199737335 --- /dev/null +++ b/tests/flytekit/unit/core/test_checkpoint.py @@ -0,0 +1,117 @@ +from pathlib import Path + +import pytest + +import flytekit +from flytekit.core.checkpointer import SyncCheckpoint + + +def test_sync_checkpoint_write(tmpdir): + td_path = Path(tmpdir) + cp = SyncCheckpoint(checkpoint_dest=tmpdir) + assert cp.read() is None + assert cp.restore() is None + dst_path = td_path.joinpath(SyncCheckpoint.TMP_DST_PATH) + assert not dst_path.exists() + cp.write(b"bytes") + assert dst_path.exists() + + +def test_sync_checkpoint_save_file(tmpdir): + td_path = Path(tmpdir) + cp = SyncCheckpoint(checkpoint_dest=tmpdir) + dst_path = td_path.joinpath(SyncCheckpoint.TMP_DST_PATH) + assert not dst_path.exists() + inp = td_path.joinpath("test") + with inp.open("wb") as f: + f.write(b"blah") + with inp.open("rb") as f: + cp.save(f) + assert dst_path.exists() + + with pytest.raises(ValueError): + # Unsupported object + cp.save(SyncCheckpoint) # noqa + + +def test_sync_checkpoint_save_filepath(tmpdir): + td_path = Path(tmpdir) + cp = SyncCheckpoint(checkpoint_dest=tmpdir) + dst_path = td_path.joinpath("test") + assert not dst_path.exists() + inp = td_path.joinpath("test") + with inp.open("wb") as f: + f.write(b"blah") + cp.save(inp) + assert dst_path.exists() + + +def test_sync_checkpoint_restore(tmpdir): + td_path = Path(tmpdir) + dest = td_path.joinpath("dest") + dest.mkdir() + src = td_path.joinpath("src") + src.mkdir() + prev = src.joinpath("prev") + p = b"prev-bytes" + with prev.open("wb") as f: + f.write(p) + cp = SyncCheckpoint(checkpoint_dest=str(dest), checkpoint_src=str(src)) + user_dest = td_path.joinpath("user_dest") + + with pytest.raises(ValueError): + cp.restore(user_dest) + + user_dest.mkdir() + assert cp.restore(user_dest) == user_dest + assert cp.restore("other_path") == user_dest + + +def test_sync_checkpoint_restore_default_path(tmpdir): + td_path = Path(tmpdir) + dest = td_path.joinpath("dest") + dest.mkdir() + src = td_path.joinpath("src") + src.mkdir() + prev = src.joinpath("prev") + p = b"prev-bytes" + with prev.open("wb") as f: + f.write(p) + cp = SyncCheckpoint(checkpoint_dest=str(dest), checkpoint_src=str(src)) + assert cp.read() == p + assert cp._prev_download_path is not None + assert cp.restore() == cp._prev_download_path + + +def test_sync_checkpoint_read_multiple_files(tmpdir): + """ + Read can only work with one file. + """ + td_path = Path(tmpdir) + dest = td_path.joinpath("dest") + dest.mkdir() + src = td_path.joinpath("src") + src.mkdir() + prev = src.joinpath("prev") + prev2 = src.joinpath("prev2") + p = b"prev-bytes" + with prev.open("wb") as f: + f.write(p) + with prev2.open("wb") as f: + f.write(p) + cp = SyncCheckpoint(checkpoint_dest=str(dest), checkpoint_src=str(src)) + + with pytest.raises(ValueError, match="Expected exactly one checkpoint - found 2"): + cp.read() + + +@flytekit.task +def t1(n: int) -> int: + ctx = flytekit.current_context() + cp = ctx.checkpoint + cp.write(bytes(n + 1)) + return n + 1 + + +def test_checkpoint_task(): + assert t1(n=5) == 6 diff --git a/tests/flytekit/unit/core/test_checkpointer.py b/tests/flytekit/unit/core/test_checkpointer.py new file mode 100644 index 0000000000..dda786545b --- /dev/null +++ b/tests/flytekit/unit/core/test_checkpointer.py @@ -0,0 +1,67 @@ +import typing +from pathlib import Path + +import py.path + +from flytekit.core.checkpointer import SyncCheckpoint + +CHECKPOINT_FILE = "cp" + + +def create_folder_write_file(tmpdir: py.path.local) -> typing.Tuple[py.path.local, py.path.local, py.path.local]: + outputs = tmpdir.mkdir("outputs") + + # Make an input test directory with one file called cp + inputs = tmpdir.mkdir("inputs") + input_file = inputs.join(CHECKPOINT_FILE) + input_file.write_text("Hello!", encoding="utf-8") + + return inputs, input_file, outputs + + +def test_sync_checkpoint_file(tmpdir: py.path.local): + inputs, input_file, outputs = create_folder_write_file(tmpdir) + cp = SyncCheckpoint(checkpoint_dest=str(outputs)) + # Lets try to restore - should not work! + assert not cp.restore("/tmp") + # Now save + cp.save(str(input_file)) + # Expect file in tmpdir + expected_dst = outputs.join(CHECKPOINT_FILE) + assert outputs.listdir() == [expected_dst] + + +def test_sync_checkpoint_reader(tmpdir: py.path.local): + inputs, input_file, outputs = create_folder_write_file(tmpdir) + cp = SyncCheckpoint(checkpoint_dest=str(outputs)) + # Lets try to restore - should not work! + assert not cp.restore("/tmp") + # Now save + with input_file.open(mode="rb") as b: + cp.save(b) + # Expect file in tmpdir + expected_dst = outputs.join(SyncCheckpoint.TMP_DST_PATH) + assert outputs.listdir() == [expected_dst] + + +def test_sync_checkpoint_folder(tmpdir: py.path.local): + inputs, input_file, outputs = create_folder_write_file(tmpdir) + cp = SyncCheckpoint(checkpoint_dest=str(outputs)) + # Lets try to restore - should not work! + assert not cp.restore("/tmp") + # Now save + cp.save(Path(str(inputs))) + # Expect file in tmpdir + expected_dst = outputs.join(CHECKPOINT_FILE) + assert outputs.listdir() == [expected_dst] + + +def test_sync_checkpoint_previous(tmpdir: py.path.local): + inputs, input_file, outputs = create_folder_write_file(tmpdir) + cp = SyncCheckpoint(checkpoint_dest=str(outputs), checkpoint_src=str(inputs)) + scratch = tmpdir.mkdir("user_scratch") + assert cp.restore(str(scratch)) == scratch + assert scratch.listdir() == [scratch.join(CHECKPOINT_FILE)] + + # ensure download is not performed again + assert cp.restore("x") == scratch diff --git a/tests/flytekit/unit/core/test_map_task.py b/tests/flytekit/unit/core/test_map_task.py index 31253ce1dd..d1f95852c1 100644 --- a/tests/flytekit/unit/core/test_map_task.py +++ b/tests/flytekit/unit/core/test_map_task.py @@ -76,6 +76,10 @@ def test_serialization(): "{{.outputPrefix}}", "--raw-output-data-prefix", "{{.rawOutputDataPrefix}}", + "--checkpoint-path", + "{{.checkpointOutputPrefix}}", + "--prev-checkpoint", + "{{.prevCheckpointPrefix}}", "--resolver", "flytekit.core.python_auto_container.default_task_resolver", "--", diff --git a/tests/flytekit/unit/core/test_python_auto_container.py b/tests/flytekit/unit/core/test_python_auto_container.py index edd3140554..4394710dd5 100644 --- a/tests/flytekit/unit/core/test_python_auto_container.py +++ b/tests/flytekit/unit/core/test_python_auto_container.py @@ -1,7 +1,9 @@ +from typing import Any + import pytest -from flytekit.core.context_manager import Image, ImageConfig -from flytekit.core.python_auto_container import get_registerable_container_image +from flytekit.core.context_manager import Image, ImageConfig, SerializationSettings +from flytekit.core.python_auto_container import PythonAutoContainerTask, get_registerable_container_image @pytest.fixture @@ -10,7 +12,46 @@ def default_image_config(): return ImageConfig(default_image=default_image) +@pytest.fixture +def default_serialization_settings(default_image_config): + return SerializationSettings( + project="p", domain="d", version="v", image_config=default_image_config, env={"FOO": "bar"} + ) + + def test_image_name_interpolation(default_image_config): img_to_interpolate = "{{.image.default.fqn}}:{{.image.default.version}}-special" img = get_registerable_container_image(img=img_to_interpolate, cfg=default_image_config) assert img == "docker.io/xyz:some-git-hash-special" + + +class DummyAutoContainerTask(PythonAutoContainerTask): + def execute(self, **kwargs) -> Any: + pass + + +task = DummyAutoContainerTask(name="x", task_config=None, task_type="t") + + +def test_default_command(default_serialization_settings): + cmd = task.get_default_command(settings=default_serialization_settings) + assert cmd == [ + "pyflyte-execute", + "--inputs", + "{{.input}}", + "--output-prefix", + "{{.outputPrefix}}", + "--raw-output-data-prefix", + "{{.rawOutputDataPrefix}}", + "--checkpoint-path", + "{{.checkpointOutputPrefix}}", + "--prev-checkpoint", + "{{.prevCheckpointPrefix}}", + "--resolver", + "flytekit.core.python_auto_container.default_task_resolver", + "--", + "task-module", + "test_python_auto_container", + "task-name", + "task", + ] From 668bf75701c80287a2852eb203248dd0c69539af Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Tue, 25 Jan 2022 13:10:15 +0800 Subject: [PATCH 068/128] Support reading subset column (#822) * Support StructuredDatasetDecoder read subset column Signed-off-by: Kevin Su * Added tests Signed-off-by: Kevin Su * Fixed tests Signed-off-by: Kevin Su * Fixed typo Signed-off-by: Kevin Su * Updated tests Signed-off-by: Kevin Su * [pr into #822] (#827) Signed-off-by: Yee Hing Tong Signed-off-by: Kevin Su * [pr into #822] Final update to structured dataset column subsetting (#828) Signed-off-by: Yee Hing Tong Co-authored-by: Yee Hing Tong --- dev-requirements.in | 2 + dev-requirements.txt | 70 +++- flytekit/__init__.py | 6 +- flytekit/core/base_task.py | 4 +- flytekit/models/types.py | 4 + flytekit/types/structured/basic_dfs.py | 11 +- flytekit/types/structured/bigquery.py | 12 +- .../types/structured/structured_dataset.py | 335 ++++++++++++------ setup.py | 1 + .../unit/core/test_structured_dataset.py | 111 ++++-- tests/flytekit/unit/core/test_type_engine.py | 28 +- tests/flytekit/unit/core/test_workflows.py | 37 +- .../types/structured_dataset/test_bigquery.py | 54 +++ .../test_structured_dataset_workflow.py | 28 +- 14 files changed, 534 insertions(+), 169 deletions(-) create mode 100644 tests/flytekit/unit/types/structured_dataset/test_bigquery.py rename tests/flytekit/unit/{type_engines => types}/structured_dataset/test_structured_dataset_workflow.py (85%) diff --git a/dev-requirements.in b/dev-requirements.in index c8eb2c0601..0c29d2fe8a 100644 --- a/dev-requirements.in +++ b/dev-requirements.in @@ -8,3 +8,5 @@ pytest mypy pre-commit codespell +google-cloud-bigquery +google-cloud-bigquery-storage diff --git a/dev-requirements.txt b/dev-requirements.txt index 4d5c69c79c..7c88c15106 100644 --- a/dev-requirements.txt +++ b/dev-requirements.txt @@ -2,7 +2,7 @@ # This file is autogenerated by pip-compile with python 3.9 # To update, run: # -# make dev-requirements.txt +# pip-compile dev-requirements.in # -e file:.#egg=flytekit # via @@ -24,11 +24,12 @@ bcrypt==3.2.0 # via # -c requirements.txt # paramiko - # secretstorage binaryornot==0.4.4 # via # -c requirements.txt # cookiecutter +cachetools==4.2.4 + # via google-auth certifi==2021.10.8 # via # -c requirements.txt @@ -122,10 +123,38 @@ flyteidl==0.21.17 # via # -c requirements.txt # flytekit +google-api-core[grpc]==2.4.0 + # via + # google-cloud-bigquery + # google-cloud-bigquery-storage + # google-cloud-core +google-auth==2.3.3 + # via + # google-api-core + # google-cloud-core +google-cloud-bigquery==2.32.0 + # via -r dev-requirements.in +google-cloud-bigquery-storage==2.11.0 + # via -r dev-requirements.in +google-cloud-core==2.2.2 + # via google-cloud-bigquery +google-crc32c==1.3.0 + # via google-resumable-media +google-resumable-media==2.1.0 + # via google-cloud-bigquery +googleapis-common-protos==1.54.0 + # via + # google-api-core + # grpcio-status grpcio==1.43.0 # via # -c requirements.txt # flytekit + # google-api-core + # google-cloud-bigquery + # grpcio-status +grpcio-status==1.43.0 + # via google-api-core identify==2.4.0 # via pre-commit idna==3.3 @@ -138,11 +167,6 @@ importlib-metadata==4.10.0 # keyring iniconfig==1.1.1 # via pytest -jeepney==0.7.1 - # via - # -c requirements.txt - # keyring - # secretstorage jinja2==3.0.3 # via # -c requirements.txt @@ -163,6 +187,8 @@ keyring==23.4.0 # via # -c requirements.txt # flytekit +libcst==0.4.0 + # via google-cloud-bigquery-storage markupsafe==2.0.1 # via # -c requirements.txt @@ -181,10 +207,6 @@ marshmallow-jsonschema==0.13.0 # via # -c requirements.txt # flytekit -secretstorage==3.3.1 - # via - # -c requirements.txt - # keyring mock==4.0.3 # via -r dev-requirements.in mypy==0.930 @@ -208,6 +230,7 @@ numpy==1.21.5 packaging==21.3 # via # -c requirements.txt + # google-cloud-bigquery # pytest pandas==1.3.5 # via @@ -229,11 +252,20 @@ poyo==0.5.0 # cookiecutter pre-commit==2.16.0 # via -r dev-requirements.in +proto-plus==1.19.8 + # via + # google-cloud-bigquery + # google-cloud-bigquery-storage protobuf==3.19.1 # via # -c requirements.txt # flyteidl # flytekit + # google-api-core + # google-cloud-bigquery + # googleapis-common-protos + # grpcio-status + # proto-plus py==1.11.0 # via # -c requirements.txt @@ -243,6 +275,12 @@ pyarrow==6.0.1 # via # -c requirements.txt # flytekit +pyasn1==0.4.8 + # via + # pyasn1-modules + # rsa +pyasn1-modules==0.2.8 + # via google-auth pycparser==2.21 # via # -c requirements.txt @@ -274,6 +312,7 @@ python-dateutil==2.8.1 # arrow # croniter # flytekit + # google-cloud-bigquery # pandas python-dotenv==0.19.2 # via docker-compose @@ -298,6 +337,7 @@ pyyaml==5.4.1 # via # -c requirements.txt # docker-compose + # libcst # pre-commit regex==2021.11.10 # via @@ -310,6 +350,8 @@ requests==2.26.0 # docker # docker-compose # flytekit + # google-api-core + # google-cloud-bigquery # responses responses==0.16.0 # via @@ -319,13 +361,15 @@ retry==0.9.2 # via # -c requirements.txt # flytekit +rsa==4.8 + # via google-auth six==1.16.0 # via # -c requirements.txt # bcrypt # cookiecutter # dockerpty - # flytekit + # google-auth # grpcio # jsonschema # pynacl @@ -359,12 +403,14 @@ tomli==1.2.3 typing-extensions==4.0.1 # via # -c requirements.txt + # libcst # mypy # typing-inspect typing-inspect==0.7.1 # via # -c requirements.txt # dataclasses-json + # libcst urllib3==1.26.7 # via # -c requirements.txt diff --git a/flytekit/__init__.py b/flytekit/__init__.py index 369316c642..fd8fcc7d9c 100644 --- a/flytekit/__init__.py +++ b/flytekit/__init__.py @@ -186,7 +186,11 @@ from flytekit.models.literals import Blob, BlobMetadata, Literal, Scalar from flytekit.models.types import LiteralType from flytekit.types import directory, file, schema -from flytekit.types.structured.structured_dataset import StructuredDataset, StructuredDatasetType +from flytekit.types.structured.structured_dataset import ( + StructuredDataset, + StructuredDatasetFormat, + StructuredDatasetType, +) __version__ = "0.0.0+develop" diff --git a/flytekit/core/base_task.py b/flytekit/core/base_task.py index c5c35373d8..cb540420db 100644 --- a/flytekit/core/base_task.py +++ b/flytekit/core/base_task.py @@ -21,7 +21,7 @@ import datetime from abc import abstractmethod from dataclasses import dataclass -from typing import Any, Dict, Generic, List, Optional, Tuple, Type, TypeVar, Union +from typing import Any, Dict, Generic, List, Optional, OrderedDict, Tuple, Type, TypeVar, Union from flytekit.core.context_manager import ( ExecutionParameters, @@ -53,7 +53,7 @@ from flytekit.models.security import SecurityContext -def kwtypes(**kwargs) -> Dict[str, Type]: +def kwtypes(**kwargs) -> OrderedDict[str, Type]: """ This is a small helper function to convert the keyword arguments to an OrderedDict of types. diff --git a/flytekit/models/types.py b/flytekit/models/types.py index a3d36de77b..50c3432838 100644 --- a/flytekit/models/types.py +++ b/flytekit/models/types.py @@ -190,6 +190,10 @@ def __init__( def columns(self) -> typing.List[DatasetColumn]: return self._columns + @columns.setter + def columns(self, value): + self._columns = value + @property def format(self) -> str: return self._format diff --git a/flytekit/types/structured/basic_dfs.py b/flytekit/types/structured/basic_dfs.py index f7a28bdc75..a65fbbeee3 100644 --- a/flytekit/types/structured/basic_dfs.py +++ b/flytekit/types/structured/basic_dfs.py @@ -2,7 +2,6 @@ import typing from typing import TypeVar -import pandas import pandas as pd import pyarrow as pa import pyarrow.parquet as pq @@ -60,6 +59,11 @@ def decode( path = flyte_value.uri local_dir = ctx.file_access.get_random_local_directory() ctx.file_access.get_data(path, local_dir, is_multipart=True) + if flyte_value.metadata.structured_dataset_type.columns: + columns = [] + for c in flyte_value.metadata.structured_dataset_type.columns: + columns.append(c.name) + return pd.read_parquet(local_dir, columns=columns) return pd.read_parquet(local_dir) @@ -94,6 +98,11 @@ def decode( path = flyte_value.uri local_dir = ctx.file_access.get_random_local_directory() ctx.file_access.get_data(path, local_dir, is_multipart=True) + if flyte_value.metadata.structured_dataset_type.columns: + columns = [] + for c in flyte_value.metadata.structured_dataset_type.columns: + columns.append(c.name) + return pq.read_table(local_dir, columns=columns) return pq.read_table(local_dir) diff --git a/flytekit/types/structured/bigquery.py b/flytekit/types/structured/bigquery.py index 33980a3e14..d9221e4110 100644 --- a/flytekit/types/structured/bigquery.py +++ b/flytekit/types/structured/bigquery.py @@ -36,10 +36,14 @@ def _read_from_bq(flyte_value: literals.StructuredDataset) -> pd.DataFrame: table = f"projects/{project_id}/datasets/{dataset_id}/tables/{table_id}" parent = "projects/{}".format(project_id) - requested_session = types.ReadSession( - table=table, - data_format=types.DataFormat.ARROW, - ) + read_options = None + if flyte_value.metadata.structured_dataset_type.columns: + columns = [] + for c in flyte_value.metadata.structured_dataset_type.columns: + columns.append(c.name) + read_options = types.ReadSession.TableReadOptions(selected_fields=columns) + + requested_session = types.ReadSession(table=table, data_format=types.DataFormat.ARROW, read_options=read_options) read_session = client.create_read_session(parent=parent, read_session=requested_session) stream = read_session.streams[0] diff --git a/flytekit/types/structured/structured_dataset.py b/flytekit/types/structured/structured_dataset.py index fb9ab4f79f..a36e50b976 100644 --- a/flytekit/types/structured/structured_dataset.py +++ b/flytekit/types/structured/structured_dataset.py @@ -9,13 +9,14 @@ from dataclasses import dataclass, field from typing import Dict, Generator, Optional, Type, Union +import pyarrow from dataclasses_json import config, dataclass_json from marshmallow import fields try: - from typing import Annotated, get_args, get_origin + from typing import Annotated, TypeAlias, get_args, get_origin except ImportError: - from typing_extensions import Annotated, get_origin, get_args + from typing_extensions import Annotated, get_origin, get_args, TypeAlias import _datetime import numpy as _np @@ -28,7 +29,7 @@ from flytekit.models import literals from flytekit.models import types as type_models from flytekit.models.literals import Literal, Scalar, StructuredDatasetMetadata -from flytekit.models.types import LiteralType, StructuredDatasetType +from flytekit.models.types import LiteralType, SchemaType, StructuredDatasetType T = typing.TypeVar("T") # StructuredDataset type or a dataframe type DF = typing.TypeVar("DF") # Dataframe type @@ -38,21 +39,25 @@ S3 = "s3" LOCAL = "/" +# For specifying the storage formats of StructuredDatasets. It's just a string, nothing fancy. +StructuredDatasetFormat: TypeAlias = str + # Storage formats -PARQUET = "parquet" +PARQUET: StructuredDatasetFormat = "parquet" @dataclass_json @dataclass class StructuredDataset(object): - uri: typing.Optional[os.PathLike] = field(default=None, metadata=config(mm_field=fields.String())) - file_format: typing.Optional[str] = field(default=PARQUET, metadata=config(mm_field=fields.String())) """ This is the user facing StructuredDataset class. Please don't confuse it with the literals.StructuredDataset class (that is just a model, a Python class representation of the protobuf). """ - FILE_FORMAT = PARQUET + uri: typing.Optional[os.PathLike] = field(default=None, metadata=config(mm_field=fields.String())) + file_format: typing.Optional[str] = field(default=PARQUET, metadata=config(mm_field=fields.String())) + + DEFAULT_FILE_FORMAT = PARQUET @classmethod def columns(cls) -> typing.Dict[str, typing.Type]: @@ -62,41 +67,6 @@ def columns(cls) -> typing.Dict[str, typing.Type]: def column_names(cls) -> typing.List[str]: return [k for k, v in cls.columns().items()] - def __class_getitem__(cls, args: typing.Union[typing.Dict[str, typing.Type], tuple]) -> Type[StructuredDataset]: - if args is None: - return cls - - format = PARQUET - if isinstance(args, tuple): - columns = args[0] - format = args[1] - else: - columns = args - - if not isinstance(columns, dict): - raise AssertionError( - f"Columns should be specified as an ordered dict " - f"of column names and their types, received {type(columns)}" - ) - - if not isinstance(format, str): - raise AssertionError(f"format should be specified as an string, received {type(format)}") - - # If nothing happened, columns and format are the default, then just use the main class - if len(columns) == 0 and format == PARQUET: - return cls - - class _TypedStructuredDataset(StructuredDataset): - # Get the type engine to see this as kind of a generic - __origin__ = StructuredDataset - FILE_FORMAT = format - - @classmethod - def columns(cls) -> typing.Dict[str, typing.Type]: - return columns - - return _TypedStructuredDataset - def __init__( self, dataframe: typing.Optional[typing.Any] = None, @@ -135,13 +105,64 @@ def all(self) -> DF: if self._dataframe_type is None: raise ValueError("No dataframe type set. Use open() to set the local dataframe type you want to use.") ctx = FlyteContextManager.current_context() - return FLYTE_DATASET_TRANSFORMER.open_as(ctx, self.literal, self._dataframe_type) + return FLYTE_DATASET_TRANSFORMER.open_as( + ctx, self.literal, self._dataframe_type, updated_metadata=self.metadata + ) def iter(self) -> Generator[DF, None, None]: if self._dataframe_type is None: raise ValueError("No dataframe type set. Use open() to set the local dataframe type you want to use.") ctx = FlyteContextManager.current_context() - return FLYTE_DATASET_TRANSFORMER.iter_as(ctx, self.literal, self._dataframe_type) + return FLYTE_DATASET_TRANSFORMER.iter_as( + ctx, self.literal, self._dataframe_type, updated_metadata=self.metadata + ) + + +def extract_cols_and_format( + t: typing.Any, +) -> typing.Tuple[Type[T], Optional[typing.OrderedDict[str, Type]], Optional[str], Optional[pa.lib.Schema]]: + """ + Helper function, just used to iterate through Annotations and extract out the following information: + - base type, if not Annotated, it will just be the type that was passed in. + - column information, as a collections.OrderedDict, + - the storage format, as a ``StructuredDatasetFormat`` (str), + - pa.lib.Schema + + If more than one of any type of thing is found, an error will be raised. + If no instances of a given type are found, then None will be returned. + + If we add more things, we should put all the returned items in a dataclass instead of just a tuple. + + :param t: The incoming type which may or may not be Annotated + :return: Tuple representing + the original type, + optional OrderedDict of columns, + optional str for the format, + optional pyarrow Schema + """ + fmt = None + ordered_dict_cols = None + pa_schema = None + if get_origin(t) is Annotated: + base_type, *annotate_args = get_args(t) + for aa in annotate_args: + if isinstance(aa, StructuredDatasetFormat): + if fmt is not None: + raise ValueError(f"A format was already specified {fmt}, cannot use {aa}") + fmt = aa + elif isinstance(aa, collections.OrderedDict): + if ordered_dict_cols is not None: + raise ValueError(f"Column information was already found {ordered_dict_cols}, cannot use {aa}") + ordered_dict_cols = aa + elif isinstance(aa, pyarrow.Schema): + if pa_schema is not None: + raise ValueError(f"Arrow schema was already found {pa_schema}, cannot use {aa}") + pa_schema = aa + return base_type, ordered_dict_cols, fmt, pa_schema + + # We return None as the format instead of parquet or something because the transformer engine may find + # a better default for the given dataframe type. + return t, ordered_dict_cols, fmt, pa_schema class StructuredDatasetEncoder(ABC): @@ -261,6 +282,25 @@ def protocol_prefix(uri: str) -> str: return LOCAL +def convert_schema_type_to_structured_dataset_type( + column_type: int, +) -> int: + if column_type == SchemaType.SchemaColumn.SchemaColumnType.INTEGER: + return type_models.SimpleType.INTEGER + if column_type == SchemaType.SchemaColumn.SchemaColumnType.FLOAT: + return type_models.SimpleType.FLOAT + if column_type == SchemaType.SchemaColumn.SchemaColumnType.STRING: + return type_models.SimpleType.STRING + if column_type == SchemaType.SchemaColumn.SchemaColumnType.DATETIME: + return type_models.SimpleType.DATETIME + if column_type == SchemaType.SchemaColumn.SchemaColumnType.DURATION: + return type_models.SimpleType.DURATION + if column_type == SchemaType.SchemaColumn.SchemaColumnType.BOOLEAN: + return type_models.SimpleType.BOOLEAN + else: + raise AssertionError(f"Unrecognized SchemaColumnType: {column_type}") + + class StructuredDatasetTransformerEngine(TypeTransformer[StructuredDataset]): """ Think of this transformer as a higher-level meta transformer that is used for all the dataframe types. @@ -368,9 +408,9 @@ def to_literal( expected: LiteralType, ) -> Literal: # Make a copy in case we need to hand off to encoders, since we can't be sure of mutations. - # Check first to see if it's even an SD type. For backwards compatibility, we may be getting a - if get_origin(python_type) is Annotated: - python_type = get_args(python_type)[0] + # Check first to see if it's even an SD type. For backwards compatibility, we may be getting a FlyteSchema + python_type, *attrs = extract_cols_and_format(python_type) + # In case it's a FlyteSchema sdt = StructuredDatasetType(format=self.DEFAULT_FORMATS.get(python_type, None)) if expected and expected.structured_dataset_type: @@ -391,7 +431,7 @@ def to_literal( # then return the original literals.StructuredDataset without invoking any encoder # # Ex. - # def t1(dataset: StructuredDataset[my_cols]) -> StructuredDataset[my_cols]: + # def t1(dataset: Annotated[StructuredDataset, my_cols]) -> Annotated[StructuredDataset, my_cols]: # return dataset if python_val._literal_sd is not None: if python_val.dataframe is not None: @@ -405,7 +445,7 @@ def to_literal( # It gets converted into a literal first, then back into a python StructuredDataset. # # Ex. - # def t2(uri: str) -> StructuredDataset[my_cols] + # def t2(uri: str) -> Annotated[StructuredDataset, my_cols] # return StructuredDataset(uri=uri) if python_val.dataframe is None: if not python_val.uri: @@ -465,61 +505,160 @@ def encode( return Literal(scalar=Scalar(structured_dataset=sd_model)) def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[T]) -> T: + """ + The only tricky thing with converting a Literal (say the output of an earlier task), to a Python value at + the start of a task execution, is the column subsetting behavior. For example, if you have, + + def t1() -> Annotated[StructuredDataset, kwtypes(col_a=int, col_b=float)]: ... + def t2(in_a: Annotated[StructuredDataset, kwtypes(col_b=float)]): ... + + where t2(in_a=t1()), when t2 does in_a.open(pd.DataFrame).all(), it should get a DataFrame + with only one column. + + +-----------------------------+-----------------------------------------+--------------------------------------+ + | | StructuredDatasetType of the incoming Literal | + +-----------------------------+-----------------------------------------+--------------------------------------+ + | StructuredDatasetType | Has columns defined | [] columns or None | + | of currently running task | | | + +=============================+=========================================+======================================+ + | Has columns | The StructuredDatasetType passed to the decoder will have the columns | + | defined | as defined by the type annotation of the currently running task. | + | | | + | | Decoders **should** then subset the incoming data to the columns requested. | + | | | + +-----------------------------+-----------------------------------------+--------------------------------------+ + | [] columns or None | StructuredDatasetType passed to decoder | StructuredDatasetType passed to the | + | | will have the columns from the incoming | decoder will have an empty list of | + | | Literal. This is the scenario where | columns. | + | | the Literal returned by the running | | + | | task will have more information than | | + | | the running task's signature. | | + +-----------------------------+-----------------------------------------+--------------------------------------+ + """ + # Detect annotations and extract out all the relevant information that the user might supply + expected_python_type, column_dict, storage_fmt, pa_schema = extract_cols_and_format(expected_python_type) + # The literal that we get in might be an old FlyteSchema. - # We'll continue to support this for the time being. - if get_origin(expected_python_type) is Annotated: - expected_python_type = get_args(expected_python_type)[0] + # We'll continue to support this for the time being. There is some duplicated logic here but let's + # keep it copy/pasted for clarity if lv.scalar.schema is not None: - sd = StructuredDataset() + schema_columns = lv.scalar.schema.type.columns + + # See the repeated logic below for comments + if column_dict is None or len(column_dict) == 0: + final_dataset_columns = [] + if schema_columns is not None and schema_columns != []: + for c in schema_columns: + final_dataset_columns.append( + StructuredDatasetType.DatasetColumn( + name=c.name, + literal_type=LiteralType( + simple=convert_schema_type_to_structured_dataset_type(c.type), + ), + ) + ) + # Dataframe will always be serialized to parquet file by FlyteSchema transformer + new_sdt = StructuredDatasetType(columns=final_dataset_columns, format=PARQUET) + else: + final_dataset_columns = self._convert_ordered_dict_of_columns_to_list(column_dict) + # Dataframe will always be serialized to parquet file by FlyteSchema transformer + new_sdt = StructuredDatasetType(columns=final_dataset_columns, format=PARQUET) + + metad = literals.StructuredDatasetMetadata(structured_dataset_type=new_sdt) sd_literal = literals.StructuredDataset( uri=lv.scalar.schema.uri, - metadata=literals.StructuredDatasetMetadata( - # Dataframe will always be serialized to parquet file by FlyteSchema transformer - structured_dataset_type=StructuredDatasetType(format=PARQUET) - ), + metadata=metad, ) - sd._literal_sd = sd_literal + if issubclass(expected_python_type, StructuredDataset): + sd = StructuredDataset(dataframe=None, metadata=metad) + sd._literal_sd = sd_literal return sd else: return self.open_as(ctx, sd_literal, df_type=expected_python_type) - # Either a StructuredDataset type or some dataframe type. + # Start handling for StructuredDataset scalars, first look at the columns + incoming_columns = lv.scalar.structured_dataset.metadata.structured_dataset_type.columns + + # If the incoming literal, also doesn't have columns, then we just have an empty list, so initialize here + final_dataset_columns = [] + # If the current running task's input does not have columns defined, or has an empty list of columns + if column_dict is None or len(column_dict) == 0: + # but if it does, then we just copy it over + if incoming_columns is not None and incoming_columns != []: + for c in incoming_columns: + final_dataset_columns.append(c) + # If the current running task's input does have columns defined + else: + final_dataset_columns = self._convert_ordered_dict_of_columns_to_list(column_dict) + + new_sdt = StructuredDatasetType( + columns=final_dataset_columns, + format=lv.scalar.structured_dataset.metadata.structured_dataset_type.format, + external_schema_type=lv.scalar.structured_dataset.metadata.structured_dataset_type.external_schema_type, + external_schema_bytes=lv.scalar.structured_dataset.metadata.structured_dataset_type.external_schema_bytes, + ) + metad = StructuredDatasetMetadata(structured_dataset_type=new_sdt) + + # A StructuredDataset type, for example + # t1(input_a: StructuredDataset) # or + # t1(input_a: Annotated[StructuredDataset, my_cols]) if issubclass(expected_python_type, StructuredDataset): - # Just save the literal for now. If in the future we find that we need the StructuredDataset type hint - # type also, we can add it. sd = expected_python_type( dataframe=None, - # Specifying these two are just done for completeness. Kind of waste since - # we're saving the whole incoming literal to _literal_sd. - metadata=lv.scalar.structured_dataset.metadata, + # Note here that the type being passed in + metadata=metad, ) sd._literal_sd = lv.scalar.structured_dataset return sd # If the requested type was not a StructuredDataset, then it means it was a plain dataframe type, which means # we should do the opening/downloading and whatever else it might entail right now. No iteration option here. - return self.open_as(ctx, lv.scalar.structured_dataset, df_type=expected_python_type) + return self.open_as(ctx, lv.scalar.structured_dataset, df_type=expected_python_type, updated_metadata=metad) - def open_as(self, ctx: FlyteContext, sd: literals.StructuredDataset, df_type: Type[DF]) -> DF: + def open_as( + self, + ctx: FlyteContext, + sd: literals.StructuredDataset, + df_type: Type[DF], + updated_metadata: Optional[StructuredDatasetMetadata] = None, + ) -> DF: + """ + + :param ctx: + :param sd: + :param df_type: + :param meta: New metadata type, since it might be different from the metadata in the literal. + :return: + """ protocol = protocol_prefix(sd.uri) decoder = self.get_decoder(df_type, protocol, sd.metadata.structured_dataset_type.format) + # todo: revisit this, we probably should add a new field to the decoder interface + if updated_metadata: + sd._metadata = updated_metadata result = decoder.decode(ctx, sd) if isinstance(result, types.GeneratorType): raise ValueError(f"Decoder {decoder} returned iterator {result} but whole value requested from {sd}") return result def iter_as( - self, ctx: FlyteContext, sd: literals.StructuredDataset, df_type: Type[DF] + self, + ctx: FlyteContext, + sd: literals.StructuredDataset, + df_type: Type[DF], + updated_metadata: Optional[StructuredDatasetMetadata] = None, ) -> Generator[DF, None, None]: protocol = protocol_prefix(sd.uri) decoder = self.DECODERS[df_type][protocol][sd.metadata.structured_dataset_type.format] + # todo: revisit this, should we add a new field to the decoder interface + if updated_metadata: + sd._metadata = updated_metadata result = decoder.decode(ctx, sd) if not isinstance(result, types.GeneratorType): raise ValueError(f"Decoder {decoder} didn't return iterator {result} but should have from {sd}") return result - def _get_dataset_column_literal_type(self, t: Type): + def _get_dataset_column_literal_type(self, t: Type) -> type_models.LiteralType: if t in self._SUPPORTED_TYPES: return self._SUPPORTED_TYPES[t] if hasattr(t, "__origin__") and t.__origin__ == list: @@ -528,39 +667,37 @@ def _get_dataset_column_literal_type(self, t: Type): return type_models.LiteralType(map_value_type=self._get_dataset_column_literal_type(t.__args__[1])) raise AssertionError(f"type {t} is currently not supported by StructuredDataset") - def _get_dataset_type(self, t: typing.Union[Type[StructuredDataset], typing.Any]) -> StructuredDatasetType: + def _convert_ordered_dict_of_columns_to_list( + self, column_map: typing.OrderedDict[str, Type] + ) -> typing.List[StructuredDatasetType.DatasetColumn]: converted_cols: typing.List[StructuredDatasetType.DatasetColumn] = [] - # Handle different kinds of annotation - # my_cols = kwtypes(x=int, y=str) - # 1. Fill in format correctly by checking for typing.annotated. For example, Annotated[pd.Dataframe, my_cols] - if get_origin(t) is Annotated: - _, *hint_args = get_args(t) - if type(hint_args[0]) is collections.OrderedDict: - for k, v in hint_args[0].items(): - lt = self._get_dataset_column_literal_type(v) - converted_cols.append(StructuredDatasetType.DatasetColumn(name=k, literal_type=lt)) - return StructuredDatasetType(columns=converted_cols, format=PARQUET) - # 3. Fill in external schema type and bytes by checking for typing.annotated metadata. - # For example, Annotated[pd.Dataframe, pa.schema([("col1", pa.int32()), ("col2", pa.string())])] - elif type(hint_args[0]) is pa.lib.Schema: - return StructuredDatasetType( - format=PARQUET, - external_schema_type="arrow", - external_schema_bytes=typing.cast(pa.lib.Schema, hint_args[0]).to_string().encode(), - ) - raise ValueError(f"Unrecognized Annotated type for StructuredDataset {t}") - - # 2. Fill in columns by checking for StructuredDataset metadata. For example, StructuredDataset[my_cols, parquet] - elif issubclass(t, StructuredDataset): - for k, v in t.columns().items(): - lt = self._get_dataset_column_literal_type(v) - converted_cols.append(StructuredDatasetType.DatasetColumn(name=k, literal_type=lt)) - return StructuredDatasetType(columns=converted_cols, format=t.FILE_FORMAT) + if column_map is None or len(column_map) == 0: + return converted_cols + for k, v in column_map.items(): + lt = self._get_dataset_column_literal_type(v) + converted_cols.append(StructuredDatasetType.DatasetColumn(name=k, literal_type=lt)) + return converted_cols - # 3. pd.Dataframe - else: - fmt = self.DEFAULT_FORMATS.get(t, PARQUET) - return StructuredDatasetType(columns=converted_cols, format=fmt) + def _get_dataset_type(self, t: typing.Union[Type[StructuredDataset], typing.Any]) -> StructuredDatasetType: + original_python_type, column_map, storage_format, pa_schema = extract_cols_and_format(t) + + # Get the column information + converted_cols = self._convert_ordered_dict_of_columns_to_list(column_map) + + # Get the format + default_format = ( + original_python_type.DEFAULT_FILE_FORMAT + if issubclass(original_python_type, StructuredDataset) + else self.DEFAULT_FORMATS.get(original_python_type, PARQUET) + ) + fmt = storage_format or default_format + + return StructuredDatasetType( + columns=converted_cols, + format=fmt, + external_schema_type="arrow" if pa_schema else None, + external_schema_bytes=typing.cast(pa.lib.Schema, pa_schema).to_string().encode() if pa_schema else None, + ) def get_literal_type(self, t: typing.Union[Type[StructuredDataset], typing.Any]) -> LiteralType: """ diff --git a/setup.py b/setup.py index ae028038bb..740a1d81e5 100644 --- a/setup.py +++ b/setup.py @@ -61,6 +61,7 @@ "natsort>=7.0.1", "docker-image-py>=0.1.10", "singledispatchmethod; python_version < '3.8.0'", + "typing_extensions", "docstring-parser>=0.9.0", "diskcache>=5.2.1", "checksumdir>=1.2.0", diff --git a/tests/flytekit/unit/core/test_structured_dataset.py b/tests/flytekit/unit/core/test_structured_dataset.py index cb8ad85071..4e4309e292 100644 --- a/tests/flytekit/unit/core/test_structured_dataset.py +++ b/tests/flytekit/unit/core/test_structured_dataset.py @@ -7,7 +7,7 @@ from flytekit.core.type_engine import TypeEngine from flytekit.models import literals from flytekit.models.literals import StructuredDatasetMetadata -from flytekit.models.types import SimpleType, StructuredDatasetType +from flytekit.models.types import SchemaType, SimpleType, StructuredDatasetType try: from typing import Annotated @@ -24,6 +24,8 @@ StructuredDataset, StructuredDatasetDecoder, StructuredDatasetEncoder, + convert_schema_type_to_structured_dataset_type, + extract_cols_and_format, protocol_prefix, ) @@ -47,7 +49,7 @@ def test_protocol(): def generate_pandas() -> pd.DataFrame: - return pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]}) + return pd.DataFrame({"name": ["Tom", "Joseph"], "age": [20, 22]}) def test_types_pandas(): @@ -58,6 +60,21 @@ def test_types_pandas(): assert lt.structured_dataset_type.columns == [] +def test_annotate_extraction(): + xyz = Annotated[pd.DataFrame, "myformat"] + a, b, c, d = extract_cols_and_format(xyz) + assert a is pd.DataFrame + assert b is None + assert c == "myformat" + assert d is None + + a, b, c, d = extract_cols_and_format(pd.DataFrame) + assert a is pd.DataFrame + assert b is None + assert c is None + assert d is None + + def test_types_annotated(): pt = Annotated[pd.DataFrame, my_cols] lt = TypeEngine.to_literal_type(pt) @@ -69,7 +86,7 @@ def test_types_annotated(): assert lt.structured_dataset_type.columns[2].literal_type.simple == SimpleType.INTEGER assert lt.structured_dataset_type.columns[3].literal_type.simple == SimpleType.STRING - pt = Annotated[pd.DataFrame, arrow_schema] + pt = Annotated[pd.DataFrame, PARQUET, arrow_schema] lt = TypeEngine.to_literal_type(pt) assert lt.structured_dataset_type.external_schema_type == "arrow" assert "some_string" in str(lt.structured_dataset_type.external_schema_bytes) @@ -78,27 +95,22 @@ def test_types_annotated(): with pytest.raises(AssertionError, match="type None is currently not supported by StructuredDataset"): TypeEngine.to_literal_type(pt) - pt = Annotated[pd.DataFrame, None] - with pytest.raises(ValueError, match="Unrecognized Annotated type for StructuredDataset"): - TypeEngine.to_literal_type(pt) - def test_types_sd(): pt = StructuredDataset lt = TypeEngine.to_literal_type(pt) assert lt.structured_dataset_type is not None - pt = StructuredDataset[my_cols] + pt = Annotated[StructuredDataset, my_cols] lt = TypeEngine.to_literal_type(pt) assert len(lt.structured_dataset_type.columns) == 4 - pt = StructuredDataset[my_cols, "csv"] + pt = Annotated[StructuredDataset, my_cols, "csv"] lt = TypeEngine.to_literal_type(pt) assert len(lt.structured_dataset_type.columns) == 4 assert lt.structured_dataset_type.format == "csv" - pt = StructuredDataset[{}, "csv"] - assert pt.FILE_FORMAT == "csv" + pt = Annotated[StructuredDataset, {}, "csv"] lt = TypeEngine.to_literal_type(pt) assert len(lt.structured_dataset_type.columns) == 0 assert lt.structured_dataset_type.format == "csv" @@ -149,7 +161,7 @@ def test_to_literal(): sd_with_uri = StructuredDataset(uri="s3://some/extant/df.parquet") - lt = TypeEngine.to_literal_type(StructuredDataset[{}, "new-df-format"]) + lt = TypeEngine.to_literal_type(Annotated[StructuredDataset, {}, "new-df-format"]) lit = FLYTE_DATASET_TRANSFORMER.to_literal(ctx, sd_with_uri, python_type=StructuredDataset, expected=lt) assert lit.scalar.structured_dataset.uri == "s3://some/extant/df.parquet" assert lit.scalar.structured_dataset.metadata.structured_dataset_type.format == "new-df-format" @@ -241,13 +253,72 @@ def decode( sd.open(pd.DataFrame).iter() -def test_class_getitem(): - assert StructuredDataset[None] == StructuredDataset - assert StructuredDataset[{}] == StructuredDataset - assert StructuredDataset[{"a": int}].FILE_FORMAT == StructuredDataset[{"a": int}, PARQUET].FILE_FORMAT +def test_convert_schema_type_to_structured_dataset_type(): + schema_ct = SchemaType.SchemaColumn.SchemaColumnType + assert convert_schema_type_to_structured_dataset_type(schema_ct.INTEGER) == SimpleType.INTEGER + assert convert_schema_type_to_structured_dataset_type(schema_ct.FLOAT) == SimpleType.FLOAT + assert convert_schema_type_to_structured_dataset_type(schema_ct.STRING) == SimpleType.STRING + assert convert_schema_type_to_structured_dataset_type(schema_ct.DATETIME) == SimpleType.DATETIME + assert convert_schema_type_to_structured_dataset_type(schema_ct.DURATION) == SimpleType.DURATION + assert convert_schema_type_to_structured_dataset_type(schema_ct.BOOLEAN) == SimpleType.BOOLEAN + with pytest.raises(AssertionError, match="Unrecognized SchemaColumnType"): + convert_schema_type_to_structured_dataset_type(int) + + +def test_to_python_value_with_incoming_columns(): + # make a literal with a type that has two columns + original_type = Annotated[pd.DataFrame, kwtypes(name=str, age=int)] + ctx = FlyteContextManager.current_context() + lt = TypeEngine.to_literal_type(original_type) + df = generate_pandas() + lit = FLYTE_DATASET_TRANSFORMER.to_literal(ctx, df, python_type=original_type, expected=lt) + assert len(lit.scalar.structured_dataset.metadata.structured_dataset_type.columns) == 2 + + # declare a new type that only has one column + # get the dataframe, make sure it has the column that was asked for. + subset_sd_type = Annotated[StructuredDataset, kwtypes(age=int)] + sd = FLYTE_DATASET_TRANSFORMER.to_python_value(ctx, lit, subset_sd_type) + assert sd.metadata.structured_dataset_type.columns[0].name == "age" + sub_df = sd.open(pd.DataFrame).all() + assert sub_df.shape[1] == 1 - with pytest.raises(AssertionError, match="Columns should be specified as an ordered dict"): - StructuredDataset[int] + # check when columns are not specified, should pull both and add column information. + sd = FLYTE_DATASET_TRANSFORMER.to_python_value(ctx, lit, StructuredDataset) + assert sd.metadata.structured_dataset_type.columns[0].name == "age" - with pytest.raises(AssertionError, match="format should be specified as an string"): - StructuredDataset[{"a": int}, 123] + # should also work if subset type is just an annotated pd.DataFrame + subset_pd_type = Annotated[pd.DataFrame, kwtypes(age=int)] + sub_df = FLYTE_DATASET_TRANSFORMER.to_python_value(ctx, lit, subset_pd_type) + assert sub_df.shape[1] == 1 + + +def test_to_python_value_without_incoming_columns(): + # make a literal with a type with no columns + ctx = FlyteContextManager.current_context() + lt = TypeEngine.to_literal_type(pd.DataFrame) + df = generate_pandas() + lit = FLYTE_DATASET_TRANSFORMER.to_literal(ctx, df, python_type=pd.DataFrame, expected=lt) + assert len(lit.scalar.structured_dataset.metadata.structured_dataset_type.columns) == 0 + + # declare a new type that only has one column + # get the dataframe, make sure it has the column that was asked for. + subset_sd_type = Annotated[StructuredDataset, kwtypes(age=int)] + sd = FLYTE_DATASET_TRANSFORMER.to_python_value(ctx, lit, subset_sd_type) + assert sd.metadata.structured_dataset_type.columns[0].name == "age" + sub_df = sd.open(pd.DataFrame).all() + assert sub_df.shape[1] == 1 + + # check when columns are not specified, should pull both and add column information. + # todo: see the todos in the open_as, and iter_as functions in StructuredDatasetTransformerEngine + # we have to recreate the literal because the test case above filled in the metadata + lit = FLYTE_DATASET_TRANSFORMER.to_literal(ctx, df, python_type=pd.DataFrame, expected=lt) + sd = FLYTE_DATASET_TRANSFORMER.to_python_value(ctx, lit, StructuredDataset) + assert sd.metadata.structured_dataset_type.columns == [] + sub_df = sd.open(pd.DataFrame).all() + assert sub_df.shape[1] == 2 + + # should also work if subset type is just an annotated pd.DataFrame + lit = FLYTE_DATASET_TRANSFORMER.to_literal(ctx, df, python_type=pd.DataFrame, expected=lt) + subset_pd_type = Annotated[pd.DataFrame, kwtypes(age=int)] + sub_df = FLYTE_DATASET_TRANSFORMER.to_python_value(ctx, lit, subset_pd_type) + assert sub_df.shape[1] == 1 diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 04ea6a0bdd..a046b58ee9 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -48,6 +48,11 @@ from flytekit.types.schema import FlyteSchema from flytekit.types.structured.structured_dataset import StructuredDataset +try: + from typing import Annotated +except ImportError: + from typing_extensions import Annotated + T = typing.TypeVar("T") @@ -641,12 +646,14 @@ def test_structured_dataset_type(): name = "Name" age = "Age" data = {name: ["Tom", "Joseph"], age: [20, 22]} + superset_cols = kwtypes(Name=str, Age=int) + subset_cols = kwtypes(Name=str) df = pd.DataFrame(data) from flytekit.types.structured.structured_dataset import StructuredDataset, StructuredDatasetTransformerEngine tf = StructuredDatasetTransformerEngine() - lt = tf.get_literal_type(StructuredDataset[{name: str, age: int}, "parquet"]) + lt = tf.get_literal_type(Annotated[StructuredDataset, superset_cols, "parquet"]) assert lt.structured_dataset_type is not None ctx = FlyteContextManager.current_context() @@ -659,6 +666,25 @@ def test_structured_dataset_type(): assert_frame_equal(df, v1) assert_frame_equal(df, v2.to_pandas()) + subset_lt = tf.get_literal_type(Annotated[StructuredDataset, subset_cols, "parquet"]) + assert subset_lt.structured_dataset_type is not None + + subset_lv = tf.to_literal(ctx, df, pd.DataFrame, subset_lt) + assert "/tmp/flyte" in subset_lv.scalar.structured_dataset.uri + v1 = tf.to_python_value(ctx, subset_lv, pd.DataFrame) + v2 = tf.to_python_value(ctx, subset_lv, pa.Table) + subset_data = pd.DataFrame({name: ["Tom", "Joseph"]}) + assert_frame_equal(subset_data, v1) + assert_frame_equal(subset_data, v2.to_pandas()) + + empty_lt = tf.get_literal_type(Annotated[StructuredDataset, "parquet"]) + assert empty_lt.structured_dataset_type is not None + empty_lv = tf.to_literal(ctx, df, pd.DataFrame, empty_lt) + v1 = tf.to_python_value(ctx, empty_lv, pd.DataFrame) + v2 = tf.to_python_value(ctx, empty_lv, pa.Table) + assert_frame_equal(df, v1) + assert_frame_equal(df, v2.to_pandas()) + def test_enum_type(): t = TypeEngine.to_literal_type(Color) diff --git a/tests/flytekit/unit/core/test_workflows.py b/tests/flytekit/unit/core/test_workflows.py index ef435393ef..c199a474de 100644 --- a/tests/flytekit/unit/core/test_workflows.py +++ b/tests/flytekit/unit/core/test_workflows.py @@ -267,55 +267,60 @@ def test_wf_docstring(): assert model_wf.template.interface.inputs["a"].description == "input a" -my_cols = kwtypes(y=int, z=int) -pd_df = pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]}) +superset_cols = kwtypes(Name=str, Age=int, Height=int) +subset_cols = kwtypes(Name=str) +superset_df = pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22], "Height": [160, 178]}) +subset_df = pd.DataFrame({"Name": ["Tom", "Joseph"]}) @task -def t1() -> Annotated[pd.DataFrame, my_cols]: - return pd_df +def t1() -> Annotated[pd.DataFrame, superset_cols]: + return superset_df @task -def t2(df: Annotated[pd.DataFrame, my_cols]) -> Annotated[pd.DataFrame, my_cols]: +def t2(df: Annotated[pd.DataFrame, subset_cols]) -> Annotated[pd.DataFrame, subset_cols]: return df @task -def t3(df: FlyteSchema[my_cols]) -> FlyteSchema[my_cols]: +def t3(df: FlyteSchema[superset_cols]) -> FlyteSchema[superset_cols]: return df @task -def t4() -> FlyteSchema[my_cols]: - return pd_df +def t4() -> FlyteSchema[superset_cols]: + return superset_df @task -def t5(sd: StructuredDataset[my_cols]) -> Annotated[pd.DataFrame, my_cols]: +def t5(sd: Annotated[StructuredDataset, subset_cols]) -> Annotated[pd.DataFrame, subset_cols]: return sd.open(pd.DataFrame).all() @workflow -def sd_wf() -> Annotated[pd.DataFrame, my_cols]: +def sd_wf() -> Annotated[pd.DataFrame, subset_cols]: + # StructuredDataset -> StructuredDataset df = t1() return t2(df=df) @workflow def sd_to_schema_wf() -> pd.DataFrame: + # StructuredDataset -> schema df = t1() return t3(df=df) @workflow -def schema_to_sd_wf() -> pd.DataFrame: +def schema_to_sd_wf() -> (pd.DataFrame, pd.DataFrame): + # schema -> StructuredDataset df = t4() - t2(df=df) - return t5(sd=df) + return t2(df=df), t5(sd=df) def test_structured_dataset_wf(): - assert_frame_equal(sd_wf(), pd_df) - assert_frame_equal(sd_to_schema_wf(), pd_df) - assert_frame_equal(schema_to_sd_wf(), pd_df) + assert_frame_equal(sd_wf(), subset_df) + assert_frame_equal(sd_to_schema_wf(), superset_df) + assert_frame_equal(schema_to_sd_wf()[0], subset_df) + assert_frame_equal(schema_to_sd_wf()[1], subset_df) diff --git a/tests/flytekit/unit/types/structured_dataset/test_bigquery.py b/tests/flytekit/unit/types/structured_dataset/test_bigquery.py new file mode 100644 index 0000000000..4801568d5e --- /dev/null +++ b/tests/flytekit/unit/types/structured_dataset/test_bigquery.py @@ -0,0 +1,54 @@ +import mock +import pandas as pd + +from flytekit import StructuredDataset, kwtypes, task, workflow + +try: + from typing import Annotated +except ImportError: + from typing_extensions import Annotated + + +pd_df = pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]}) +my_cols = kwtypes(Name=str, Age=int) + + +@task +def gen_df() -> Annotated[pd.DataFrame, my_cols, "parquet"]: + return pd_df + + +@task +def t1(df: pd.DataFrame) -> Annotated[StructuredDataset, my_cols]: + return StructuredDataset(dataframe=df, uri="bq://project:flyte.table") + + +@task +def t2(sd: Annotated[StructuredDataset, my_cols]) -> pd.DataFrame: + return sd.open(pd.DataFrame).all() + + +@workflow +def wf() -> pd.DataFrame: + df = gen_df() + sd = t1(df=df) + return t2(sd=sd) + + +@mock.patch("google.cloud.bigquery.Client") +@mock.patch("google.cloud.bigquery_storage.BigQueryReadClient") +@mock.patch("google.cloud.bigquery_storage_v1.reader.ReadRowsStream") +def test_bq_wf(mock_read_rows_stream, mock_bigquery_read_client, mock_client): + class mock_pages: + def to_dataframe(self): + return pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]}) + + class mock_rows: + pages = [mock_pages()] + + mock_client.load_table_from_dataframe.return_value = None + mock_read_rows_stream.rows.return_value = mock_rows + mock_bigquery_read_client.read_rows.return_value = mock_read_rows_stream + mock_bigquery_read_client.return_value = mock_bigquery_read_client + + assert wf().equals(pd_df) diff --git a/tests/flytekit/unit/type_engines/structured_dataset/test_structured_dataset_workflow.py b/tests/flytekit/unit/types/structured_dataset/test_structured_dataset_workflow.py similarity index 85% rename from tests/flytekit/unit/type_engines/structured_dataset/test_structured_dataset_workflow.py rename to tests/flytekit/unit/types/structured_dataset/test_structured_dataset_workflow.py index ae19956b3c..fa879ff0d9 100644 --- a/tests/flytekit/unit/type_engines/structured_dataset/test_structured_dataset_workflow.py +++ b/tests/flytekit/unit/types/structured_dataset/test_structured_dataset_workflow.py @@ -31,8 +31,8 @@ NUMPY_PATH = FlyteContextManager.current_context().file_access.get_random_local_directory() BQ_PATH = "bq://flyte-dataset:flyte.table" -my_cols = kwtypes(w=typing.Dict[str, typing.Dict[str, int]], x=typing.List[typing.List[int]], y=int, z=str) -fields = [("some_int", pa.int32()), ("some_string", pa.string())] +my_cols = kwtypes(Name=str, Age=int) +fields = [("Name", pa.string()), ("Age", pa.int32())] arrow_schema = pa.schema(fields) pd_df = pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]}) @@ -45,7 +45,7 @@ def encode( structured_dataset_type: StructuredDatasetType, ) -> literals.StructuredDataset: return literals.StructuredDataset( - uri="s3://bucket/key", metadata=StructuredDatasetMetadata(structured_dataset_type) + uri="bq://bucket/key", metadata=StructuredDatasetMetadata(structured_dataset_type) ) @@ -106,7 +106,7 @@ def t1(dataframe: pd.DataFrame) -> Annotated[pd.DataFrame, my_cols]: @task -def t1a(dataframe: pd.DataFrame) -> StructuredDataset[my_cols, PARQUET]: +def t1a(dataframe: pd.DataFrame) -> Annotated[StructuredDataset, my_cols, PARQUET]: # S3 (parquet) -> Pandas -> S3 (parquet) return StructuredDataset(dataframe=dataframe, uri=PANDAS_PATH) @@ -118,7 +118,7 @@ def t2(dataframe: pd.DataFrame) -> Annotated[pd.DataFrame, arrow_schema]: @task -def t3(dataset: StructuredDataset[my_cols]) -> StructuredDataset[my_cols]: +def t3(dataset: Annotated[StructuredDataset, my_cols]) -> Annotated[StructuredDataset, my_cols]: # s3 (parquet) -> pandas -> s3 (parquet) print(dataset.open(pd.DataFrame).all()) # In the example, we download dataset when we open it. @@ -127,39 +127,41 @@ def t3(dataset: StructuredDataset[my_cols]) -> StructuredDataset[my_cols]: @task -def t3a(dataset: StructuredDataset[my_cols]) -> StructuredDataset[my_cols]: +def t3a(dataset: Annotated[StructuredDataset, my_cols]) -> Annotated[StructuredDataset, my_cols]: # This task will not do anything - no uploading, no downloading return dataset @task -def t4(dataset: StructuredDataset[my_cols]) -> pd.DataFrame: +def t4(dataset: Annotated[StructuredDataset, my_cols]) -> pd.DataFrame: # s3 (parquet) -> pandas -> s3 (parquet) return dataset.open(pd.DataFrame).all() @task -def t5(dataframe: pd.DataFrame) -> StructuredDataset[my_cols]: +def t5(dataframe: pd.DataFrame) -> Annotated[StructuredDataset, my_cols]: # s3 (parquet) -> pandas -> bq return StructuredDataset(dataframe=dataframe, uri=BQ_PATH) @task -def t6(dataset: StructuredDataset[my_cols]) -> pd.DataFrame: +def t6(dataset: Annotated[StructuredDataset, my_cols]) -> pd.DataFrame: # bq -> pandas -> s3 (parquet) df = dataset.open(pd.DataFrame).all() return df @task -def t7(df1: pd.DataFrame, df2: pd.DataFrame) -> (StructuredDataset[my_cols], StructuredDataset[my_cols]): +def t7( + df1: pd.DataFrame, df2: pd.DataFrame +) -> (Annotated[StructuredDataset, my_cols], Annotated[StructuredDataset, my_cols]): # df1: pandas -> bq # df2: pandas -> s3 (parquet) return StructuredDataset(dataframe=df1, uri=BQ_PATH), StructuredDataset(dataframe=df2) @task -def t8(dataframe: pa.Table) -> StructuredDataset[my_cols]: +def t8(dataframe: pa.Table) -> Annotated[StructuredDataset, my_cols]: # Arrow table -> s3 (parquet) print(dataframe.columns) return StructuredDataset(dataframe=dataframe) @@ -173,13 +175,13 @@ def t8a(dataframe: pa.Table) -> pa.Table: @task -def t9(dataframe: np.ndarray) -> StructuredDataset[my_cols]: +def t9(dataframe: np.ndarray) -> Annotated[StructuredDataset, my_cols]: # numpy -> Arrow table -> s3 (parquet) return StructuredDataset(dataframe=dataframe, uri=NUMPY_PATH) @task -def t10(dataset: StructuredDataset[my_cols]) -> np.ndarray: +def t10(dataset: Annotated[StructuredDataset, my_cols]) -> np.ndarray: # s3 (parquet) -> Arrow table -> numpy np_array = dataset.open(np.ndarray).all() return np_array From 11f28c3e95b62be3e7d3a28fd55966b6235cd7f0 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Thu, 27 Jan 2022 01:43:19 +0800 Subject: [PATCH 069/128] Fix spark regression (#830) Signed-off-by: Kevin Su Signed-off-by: maximsmol --- plugins/flytekit-spark/flytekitplugins/spark/task.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/plugins/flytekit-spark/flytekitplugins/spark/task.py b/plugins/flytekit-spark/flytekitplugins/spark/task.py index c05dbc2ea0..b59fb3aed5 100644 --- a/plugins/flytekit-spark/flytekitplugins/spark/task.py +++ b/plugins/flytekit-spark/flytekitplugins/spark/task.py @@ -106,7 +106,7 @@ def pre_execute(self, user_params: ExecutionParameters) -> ExecutionParameters: ctx = FlyteContextManager.current_context() sess_builder = _pyspark.sql.SparkSession.builder.appName(f"FlyteSpark: {user_params.execution_id}") - if not (ctx.execution_state and ctx.execution_state.Mode == ExecutionState.Mode.TASK_EXECUTION): + if not (ctx.execution_state and ctx.execution_state.mode == ExecutionState.Mode.TASK_EXECUTION): # If either of above cases is not true, then we are in local execution of this task # Add system spark-conf for local/notebook based execution. spark_conf = _pyspark.SparkConf() From 79e727fad9d891aa961a397017833e6befdbedaf Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Thu, 27 Jan 2022 14:44:31 -0800 Subject: [PATCH 070/128] Update argument setting for in fast registered, dynamically generated, pod tasks (#835) Signed-off-by: Yee Hing Tong Signed-off-by: maximsmol --- flytekit/core/context_manager.py | 4 ++ flytekit/core/python_function_task.py | 20 +++--- flytekit/tools/translator.py | 14 ++--- .../flytekitplugins/pod/task.py | 2 +- plugins/flytekit-k8s-pod/tests/test_pod.py | 63 +++++++++++++++++++ tests/flytekit/unit/core/test_type_hints.py | 4 +- 6 files changed, 84 insertions(+), 23 deletions(-) diff --git a/flytekit/core/context_manager.py b/flytekit/core/context_manager.py index 70af22f1e1..0b58a6d14b 100644 --- a/flytekit/core/context_manager.py +++ b/flytekit/core/context_manager.py @@ -400,8 +400,12 @@ class FastSerializationSettings(object): """ enabled: bool = False + # This is the location that the code should be copied into. destination_dir: Optional[str] = None + # This is the zip file where the new code was uploaded to. + distribution_location: Optional[str] = None + @dataclass(frozen=True) class SerializationSettings(object): diff --git a/flytekit/core/python_function_task.py b/flytekit/core/python_function_task.py index fa98b7ca89..35016f9571 100644 --- a/flytekit/core/python_function_task.py +++ b/flytekit/core/python_function_task.py @@ -238,18 +238,6 @@ def compile_into_workflow( "Compilation for a dynamic workflow called in fast execution mode but no additional code " "distribution could be retrieved" ) - logger.warn(f"ctx.execution_state.additional_context {ctx.execution_state.additional_context}") - for task_template in tts: - sanitized_args = [] - for arg in task_template.container.args: - if arg == "{{ .remote_package_path }}": - sanitized_args.append(ctx.execution_state.additional_context.get("dynamic_addl_distro")) - elif arg == "{{ .dest_dir }}": - sanitized_args.append(ctx.execution_state.additional_context.get("dynamic_dest_dir", ".")) - else: - sanitized_args.append(arg) - del task_template.container.args[:] - task_template.container.args.extend(sanitized_args) dj_spec = _dynamic_job.DynamicJobSpec( min_successes=len(workflow_spec.template.nodes), @@ -290,7 +278,13 @@ def dynamic_execute(self, task_function: Callable, **kwargs) -> Any: if is_fast_execution: ctx = ctx.with_serialization_settings( ctx.serialization_settings.new_builder() - .with_fast_serialization_settings(FastSerializationSettings(enabled=True)) + .with_fast_serialization_settings( + FastSerializationSettings( + enabled=True, + destination_dir=ctx.execution_state.additional_context.get("dynamic_dest_dir", "."), + distribution_location=ctx.execution_state.additional_context.get("dynamic_addl_distro"), + ) + ) .build() ) diff --git a/flytekit/tools/translator.py b/flytekit/tools/translator.py index 83d724e4ce..8c9750e8cd 100644 --- a/flytekit/tools/translator.py +++ b/flytekit/tools/translator.py @@ -68,19 +68,17 @@ def _fast_serialize_command_fn( ) -> Callable[[SerializationSettings], List[str]]: default_command = task.get_default_command(settings) - dest_dir = ( - settings.fast_serialization_settings.destination_dir if settings.fast_serialization_settings is not None else "" - ) - if dest_dir is None or dest_dir == "": - dest_dir = "{{ .dest_dir }}" - def fn(settings: SerializationSettings) -> List[str]: return [ "pyflyte-fast-execute", "--additional-distribution", - "{{ .remote_package_path }}", + settings.fast_serialization_settings.distribution_location + if settings.fast_serialization_settings and settings.fast_serialization_settings.distribution_location + else "{{ .remote_package_path }}", "--dest-dir", - dest_dir, + settings.fast_serialization_settings.destination_dir + if settings.fast_serialization_settings and settings.fast_serialization_settings.destination_dir + else "{{ .dest_dir }}", "--", *default_command, ] diff --git a/plugins/flytekit-k8s-pod/flytekitplugins/pod/task.py b/plugins/flytekit-k8s-pod/flytekitplugins/pod/task.py index 6cbf7d57fa..a4872bb321 100644 --- a/plugins/flytekit-k8s-pod/flytekitplugins/pod/task.py +++ b/plugins/flytekit-k8s-pod/flytekitplugins/pod/task.py @@ -75,7 +75,7 @@ def _serialize_pod_spec(self, settings: SerializationSettings) -> Dict[str, Any] final_containers = [] for container in containers: # In the case of the primary container, we overwrite specific container attributes with the default values - # used in an SDK runnable task. + # used in the regular Python task. if container.name == self.task_config.primary_container_name: sdk_default_container = super().get_container(settings) diff --git a/plugins/flytekit-k8s-pod/tests/test_pod.py b/plugins/flytekit-k8s-pod/tests/test_pod.py index 9677f51211..85fd2d8041 100644 --- a/plugins/flytekit-k8s-pod/tests/test_pod.py +++ b/plugins/flytekit-k8s-pod/tests/test_pod.py @@ -10,6 +10,7 @@ from flytekit import Resources, TaskMetadata, dynamic, map_task, task from flytekit.core import context_manager from flytekit.core.context_manager import FastSerializationSettings +from flytekit.core.type_engine import TypeEngine from flytekit.extend import ExecutionState, Image, ImageConfig, SerializationSettings from flytekit.tools.translator import get_serializable @@ -391,3 +392,65 @@ def simple_pod_task(i: int): "task-name", "simple_pod_task", ] + + +def test_fast(): + REQUESTS_GPU = Resources(cpu="123m", mem="234Mi", ephemeral_storage="123M", gpu="1") + LIMITS_GPU = Resources(cpu="124M", mem="235Mi", ephemeral_storage="124M", gpu="1") + + def get_minimal_pod_task_config() -> Pod: + primary_container = V1Container(name="flytetask") + pod_spec = V1PodSpec(containers=[primary_container]) + return Pod(pod_spec=pod_spec, primary_container_name="flytetask") + + @task( + task_config=get_minimal_pod_task_config(), + requests=REQUESTS_GPU, + limits=LIMITS_GPU, + ) + def pod_task_with_resources(dummy_input: str) -> str: + return dummy_input + + @dynamic(requests=REQUESTS_GPU, limits=LIMITS_GPU) + def dynamic_task_with_pod_subtask(dummy_input: str) -> str: + pod_task_with_resources(dummy_input=dummy_input) + return dummy_input + + default_img = Image(name="default", fqn="test", tag="tag") + serialization_settings = SerializationSettings( + project="project", + domain="domain", + version="version", + env={"FOO": "baz"}, + image_config=ImageConfig(default_image=default_img, images=[default_img]), + fast_serialization_settings=FastSerializationSettings(enabled=True), + ) + + with context_manager.FlyteContextManager.with_context( + context_manager.FlyteContextManager.current_context().with_serialization_settings(serialization_settings) + ) as ctx: + with context_manager.FlyteContextManager.with_context( + ctx.with_execution_state( + ctx.execution_state.with_params( + mode=ExecutionState.Mode.TASK_EXECUTION, + additional_context={ + "dynamic_addl_distro": "s3://my-s3-bucket/fast/123", + "dynamic_dest_dir": "/User/flyte/workflows", + }, + ) + ) + ) as ctx: + input_literal_map = TypeEngine.dict_to_literal_map(ctx, {"dummy_input": "hi"}) + dynamic_job_spec = dynamic_task_with_pod_subtask.dispatch_execute(ctx, input_literal_map) + # print(dynamic_job_spec) + assert len(dynamic_job_spec._nodes) == 1 + assert len(dynamic_job_spec.tasks) == 1 + args = " ".join(dynamic_job_spec.tasks[0].k8s_pod.pod_spec["containers"][0]["args"]) + assert args.startswith( + "pyflyte-fast-execute --additional-distribution s3://my-s3-bucket/fast/123 " + "--dest-dir /User/flyte/workflows" + ) + assert dynamic_job_spec.tasks[0].k8s_pod.pod_spec["containers"][0]["resources"]["limits"]["cpu"] == "124M" + assert dynamic_job_spec.tasks[0].k8s_pod.pod_spec["containers"][0]["resources"]["requests"]["gpu"] == "1" + + assert context_manager.FlyteContextManager.size() == 1 diff --git a/tests/flytekit/unit/core/test_type_hints.py b/tests/flytekit/unit/core/test_type_hints.py index 37deb6d0d7..40442c72d3 100644 --- a/tests/flytekit/unit/core/test_type_hints.py +++ b/tests/flytekit/unit/core/test_type_hints.py @@ -625,7 +625,9 @@ def my_wf(a: int) -> typing.List[str]: ) ) ) as ctx: - dynamic_job_spec = my_subwf.compile_into_workflow(ctx, my_subwf._task_function, a=5) + input_literal_map = TypeEngine.dict_to_literal_map(ctx, {"a": 5}) + + dynamic_job_spec = my_subwf.dispatch_execute(ctx, input_literal_map) assert len(dynamic_job_spec._nodes) == 5 assert len(dynamic_job_spec.tasks) == 1 args = " ".join(dynamic_job_spec.tasks[0].container.args) From 12027060c7ea0ce1de898ab6f149d9345f5bd76a Mon Sep 17 00:00:00 2001 From: Ketan Umare <16888709+kumare3@users.noreply.github.com> Date: Fri, 28 Jan 2022 16:03:36 -0800 Subject: [PATCH 071/128] `ctx` Context can be used within shell tasks - to access context vars and secrets (#832) * Adding context to a substitutable parameter in shell task Signed-off-by: Ketan Umare * Support for secrets in context Signed-off-by: Ketan Umare * addressed comments Signed-off-by: Ketan Umare Signed-off-by: maximsmol --- flytekit/core/context_manager.py | 22 +++++ flytekit/extras/tasks/shell.py | 26 ++++-- .../unit/core/test_context_manager.py | 87 ++++++++++++++++++- .../flytekit/unit/extras/tasks/test_shell.py | 72 +++++++++------ .../unit/extras/tasks/testdata/script.sh | 4 +- 5 files changed, 177 insertions(+), 34 deletions(-) diff --git a/flytekit/core/context_manager.py b/flytekit/core/context_manager.py index 0b58a6d14b..b57ff043b5 100644 --- a/flytekit/core/context_manager.py +++ b/flytekit/core/context_manager.py @@ -337,11 +337,33 @@ class SecretsManager(object): All configuration values can always be overridden by injecting an environment variable """ + class _GroupSecrets(object): + """ + This is a dummy class whose sole purpose is to support "attribute" style lookup for secrets + """ + + def __init__(self, group: str, sm: typing.Any): + self._group = group + self._sm = sm + + def __getattr__(self, item: str) -> str: + """ + Returns the secret that matches "group"."key" + the key, here is the item + """ + return self._sm.get(self._group, item) + def __init__(self): self._base_dir = str(secrets.SECRETS_DEFAULT_DIR.get()).strip() self._file_prefix = str(secrets.SECRETS_FILE_PREFIX.get()).strip() self._env_prefix = str(secrets.SECRETS_ENV_PREFIX.get()).strip() + def __getattr__(self, item: str) -> _GroupSecrets: + """ + returns a new _GroupSecrets objects, that allows all keys within this group to be looked up like attributes + """ + return self._GroupSecrets(item, self) + def get(self, group: str, key: str) -> str: """ Retrieves a secret using the resolution order -> Env followed by file. If not found raises a ValueError diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py index 63bae594ed..2686c899d8 100644 --- a/flytekit/extras/tasks/shell.py +++ b/flytekit/extras/tasks/shell.py @@ -1,4 +1,3 @@ -import collections import datetime import logging import os @@ -7,6 +6,7 @@ import typing from dataclasses import dataclass +import flytekit from flytekit.core.context_manager import ExecutionParameters from flytekit.core.interface import Interface from flytekit.core.python_function_task import PythonInstanceTask @@ -38,7 +38,15 @@ def _dummy_task_func(): return None -T = typing.TypeVar("T") +class AttrDict(dict): + """ + Convert a dictionary to an attribute style lookup. Do not use this in regular places, this is used for + namespacing inputs and outputs + """ + + def __init__(self, *args, **kwargs): + super(AttrDict, self).__init__(*args, **kwargs) + self.__dict__ = self class _PythonFStringInterpolizer: @@ -73,16 +81,22 @@ def interpolate( """ inputs = inputs or {} outputs = outputs or {} - reused_vars = inputs.keys() & outputs.keys() - if reused_vars: - raise ValueError(f"Variables {reused_vars} in Query cannot be shared between inputs and outputs.") - consolidated_args = collections.ChainMap(inputs, outputs) + inputs = AttrDict(inputs) + outputs = AttrDict(outputs) + consolidated_args = { + "inputs": inputs, + "outputs": outputs, + "ctx": flytekit.current_context(), + } try: return self._Formatter().format(tmpl, **consolidated_args) except KeyError as e: raise ValueError(f"Variable {e} in Query not found in inputs {consolidated_args.keys()}") +T = typing.TypeVar("T") + + class ShellTask(PythonInstanceTask[T]): """ """ diff --git a/tests/flytekit/unit/core/test_context_manager.py b/tests/flytekit/unit/core/test_context_manager.py index 0e7588fdc6..d237090367 100644 --- a/tests/flytekit/unit/core/test_context_manager.py +++ b/tests/flytekit/unit/core/test_context_manager.py @@ -1,4 +1,16 @@ -from flytekit.core.context_manager import ExecutionState, FlyteContext, FlyteContextManager, look_up_image_info +import os + +import py +import pytest + +from flytekit.configuration import secrets +from flytekit.core.context_manager import ( + ExecutionState, + FlyteContext, + FlyteContextManager, + SecretsManager, + look_up_image_info, +) class SampleTestClass(object): @@ -65,3 +77,76 @@ def test_additional_context(): ) ) as exec_ctx_inner: assert exec_ctx_inner.execution_state.additional_context == {1: "inner", 2: "foo", 3: "baz"} + + +def test_secrets_manager_default(): + with pytest.raises(ValueError): + sec = SecretsManager() + sec.get("group", "key") + + with pytest.raises(ValueError): + _ = sec.group.key + + +def test_secrets_manager_get_envvar(): + sec = SecretsManager() + with pytest.raises(ValueError): + sec.get_secrets_env_var("test", "") + with pytest.raises(ValueError): + sec.get_secrets_env_var("", "x") + assert sec.get_secrets_env_var("group", "test") == f"{secrets.SECRETS_ENV_PREFIX.get()}GROUP_TEST" + + +def test_secrets_manager_get_file(): + sec = SecretsManager() + with pytest.raises(ValueError): + sec.get_secrets_file("test", "") + with pytest.raises(ValueError): + sec.get_secrets_file("", "x") + assert sec.get_secrets_file("group", "test") == os.path.join( + secrets.SECRETS_DEFAULT_DIR.get(), + "group", + f"{secrets.SECRETS_FILE_PREFIX.get()}test", + ) + + +def test_secrets_manager_file(tmpdir: py.path.local): + tmp = tmpdir.mkdir("file_test").dirname + os.environ["FLYTE_SECRETS_DEFAULT_DIR"] = tmp + sec = SecretsManager() + f = os.path.join(tmp, "test") + with open(f, "w+") as w: + w.write("my-password") + + with pytest.raises(ValueError): + sec.get("test", "") + with pytest.raises(ValueError): + sec.get("", "x") + # Group dir not exists + with pytest.raises(ValueError): + sec.get("group", "test") + + g = os.path.join(tmp, "group") + os.makedirs(g) + f = os.path.join(g, "test") + with open(f, "w+") as w: + w.write("my-password") + assert sec.get("group", "test") == "my-password" + assert sec.group.test == "my-password" + del os.environ["FLYTE_SECRETS_DEFAULT_DIR"] + + +def test_secrets_manager_bad_env(): + with pytest.raises(ValueError): + os.environ["TEST"] = "value" + sec = SecretsManager() + sec.get("group", "test") + + +def test_secrets_manager_env(): + sec = SecretsManager() + os.environ[sec.get_secrets_env_var("group", "test")] = "value" + assert sec.get("group", "test") == "value" + + os.environ[sec.get_secrets_env_var(group="group", key="key")] = "value" + assert sec.get(group="group", key="key") == "value" diff --git a/tests/flytekit/unit/extras/tasks/test_shell.py b/tests/flytekit/unit/extras/tasks/test_shell.py index 2b76f7ad7f..f0ed65705e 100644 --- a/tests/flytekit/unit/extras/tasks/test_shell.py +++ b/tests/flytekit/unit/extras/tasks/test_shell.py @@ -7,6 +7,7 @@ import pytest from dataclasses_json import dataclass_json +import flytekit from flytekit import kwtypes from flytekit.extras.tasks.shell import OutputLocation, ShellTask from flytekit.types.directory import FlyteDirectory @@ -46,8 +47,8 @@ def test_input_substitution_primitive(): name="test", script=""" set -ex - cat {f} - echo "Hello World {y} on {j}" + cat {inputs.f} + echo "Hello World {inputs.y} on {inputs.j}" """, inputs=kwtypes(f=str, y=int, j=datetime.datetime), ) @@ -62,8 +63,8 @@ def test_input_substitution_files(): t = ShellTask( name="test", script=""" - cat {f} - echo "Hello World {y} on {j}" + cat {inputs.f} + echo "Hello World {inputs.y} on {inputs.j}" """, inputs=kwtypes(f=CSVFile, y=FlyteDirectory, j=datetime.datetime), ) @@ -71,15 +72,37 @@ def test_input_substitution_files(): assert t(f=test_csv, y=testdata, j=datetime.datetime(2021, 11, 10, 12, 15, 0)) is None +def test_input_substitution_files_ctx(): + sec = flytekit.current_context().secrets + envvar = sec.get_secrets_env_var("group", "key") + os.environ[envvar] = "value" + assert sec.get("group", "key") == "value" + + t = ShellTask( + name="test", + script=""" + export EXEC={ctx.execution_id} + export SECRET={ctx.secrets.group.key} + cat {inputs.f} + echo "Hello World {inputs.y} on {inputs.j}" + """, + inputs=kwtypes(f=CSVFile, y=FlyteDirectory, j=datetime.datetime), + debug=True, + ) + + assert t(f=test_csv, y=testdata, j=datetime.datetime(2021, 11, 10, 12, 15, 0)) is None + del os.environ[envvar] + + def test_input_output_substitution_files(): - script = "cat {f} > {y}" + script = "cat {inputs.f} > {outputs.y}" t = ShellTask( name="test", debug=True, script=script, inputs=kwtypes(f=CSVFile), output_locs=[ - OutputLocation(var="y", var_type=FlyteFile, location="{f}.mod"), + OutputLocation(var="y", var_type=FlyteFile, location="{inputs.f}.mod"), ], ) @@ -101,15 +124,15 @@ def test_input_output_substitution_files(): def test_input_single_output_substitution_files(): script = """ - cat {f} >> {z} - echo "Hello World {y} on {j}" + cat {inputs.f} >> {outputs.z} + echo "Hello World {inputs.y} on {inputs.j}" """ t = ShellTask( name="test", debug=True, script=script, inputs=kwtypes(f=CSVFile, y=FlyteDirectory, j=datetime.datetime), - output_locs=[OutputLocation(var="z", var_type=FlyteFile, location="{f}.pyc")], + output_locs=[OutputLocation(var="z", var_type=FlyteFile, location="{inputs.f}.pyc")], ) assert t.script == script @@ -122,14 +145,14 @@ def test_input_single_output_substitution_files(): [ ( """ - cat {missing} >> {z} - echo "Hello World {y} on {j} - output {x}" + cat {missing} >> {outputs.z} + echo "Hello World {inputs.y} on {inputs.j} - output {outputs.x}" """ ), ( """ - cat {f} {missing} >> {z} - echo "Hello World {y} on {j} - output {x}" + cat {inputs.f} {missing} >> {outputs.z} + echo "Hello World {inputs.y} on {inputs.j} - output {outputs.x}" """ ), ], @@ -141,8 +164,8 @@ def test_input_output_extra_and_missing_variables(script): script=script, inputs=kwtypes(f=CSVFile, y=FlyteDirectory, j=datetime.datetime), output_locs=[ - OutputLocation(var="x", var_type=FlyteDirectory, location="{y}"), - OutputLocation(var="z", var_type=FlyteFile, location="{f}.pyc"), + OutputLocation(var="x", var_type=FlyteDirectory, location="{inputs.y}"), + OutputLocation(var="z", var_type=FlyteFile, location="{inputs.f}.pyc"), ], ) @@ -150,22 +173,21 @@ def test_input_output_extra_and_missing_variables(script): t(f=test_csv, y=testdata, j=datetime.datetime(2021, 11, 10, 12, 15, 0)) -def test_cannot_reuse_variables_for_both_inputs_and_outputs(): +def test_reuse_variables_for_both_inputs_and_outputs(): t = ShellTask( name="test", debug=True, script=""" - cat {f} >> {y} - echo "Hello World {y} on {j}" + cat {inputs.f} >> {outputs.y} + echo "Hello World {inputs.y} on {inputs.j}" """, inputs=kwtypes(f=CSVFile, y=FlyteDirectory, j=datetime.datetime), output_locs=[ - OutputLocation(var="y", var_type=FlyteFile, location="{f}.pyc"), + OutputLocation(var="y", var_type=FlyteFile, location="{inputs.f}.pyc"), ], ) - with pytest.raises(ValueError, match="Variables {'y'} in Query"): - t(f=test_csv, y=testdata, j=datetime.datetime(2021, 11, 10, 12, 15, 0)) + t(f=test_csv, y=testdata, j=datetime.datetime(2021, 11, 10, 12, 15, 0)) def test_can_use_complex_types_for_inputs_to_f_string_template(): @@ -177,10 +199,10 @@ class InputArgs: t = ShellTask( name="test", debug=True, - script="""cat {input_args.in_file} >> {input_args.in_file}.tmp""", + script="""cat {inputs.input_args.in_file} >> {inputs.input_args.in_file}.tmp""", inputs=kwtypes(input_args=InputArgs), output_locs=[ - OutputLocation(var="x", var_type=FlyteFile, location="{input_args.in_file}.tmp"), + OutputLocation(var="x", var_type=FlyteFile, location="{inputs.input_args.in_file}.tmp"), ], ) @@ -196,8 +218,8 @@ def test_shell_script(): script_file=script_sh, inputs=kwtypes(f=CSVFile, y=FlyteDirectory, j=datetime.datetime), output_locs=[ - OutputLocation(var="x", var_type=FlyteDirectory, location="{y}"), - OutputLocation(var="z", var_type=FlyteFile, location="{f}.pyc"), + OutputLocation(var="x", var_type=FlyteDirectory, location="{inputs.y}"), + OutputLocation(var="z", var_type=FlyteFile, location="{inputs.f}.pyc"), ], ) diff --git a/tests/flytekit/unit/extras/tasks/testdata/script.sh b/tests/flytekit/unit/extras/tasks/testdata/script.sh index 1deb4c474a..f0fc6924c0 100644 --- a/tests/flytekit/unit/extras/tasks/testdata/script.sh +++ b/tests/flytekit/unit/extras/tasks/testdata/script.sh @@ -2,5 +2,5 @@ set -ex -cat "{f}" >> "{z}" -echo "Hello World {y} on {j} - output {x}" +cat "{inputs.f}" >> "{outputs.z}" +echo "Hello World {inputs.y} on {inputs.j} - output {outputs.x}" From 0f4fb9e4914ae295d6462b017c27bd9ca5004720 Mon Sep 17 00:00:00 2001 From: Ketan Umare <16888709+kumare3@users.noreply.github.com> Date: Sun, 30 Jan 2022 22:10:28 -0800 Subject: [PATCH 072/128] Expose Checkpoint as a top-level interface in flytekit (#839) Signed-off-by: Ketan Umare Signed-off-by: maximsmol --- flytekit/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/flytekit/__init__.py b/flytekit/__init__.py index fd8fcc7d9c..0c2db39db0 100644 --- a/flytekit/__init__.py +++ b/flytekit/__init__.py @@ -162,6 +162,7 @@ from flytekit.core.base_sql_task import SQLTask from flytekit.core.base_task import SecurityContext, TaskMetadata, kwtypes +from flytekit.core.checkpointer import Checkpoint from flytekit.core.condition import conditional from flytekit.core.container_task import ContainerTask from flytekit.core.context_manager import ExecutionParameters, FlyteContext, FlyteContextManager From 9b6a6cc28c87d8b03ca3e60f59368e1eca7d1ad8 Mon Sep 17 00:00:00 2001 From: bstadlbauer <11799671+bstadlbauer@users.noreply.github.com> Date: Mon, 31 Jan 2022 19:05:28 +0100 Subject: [PATCH 073/128] Parse duration field from flyteidl to `flytekit.models.execution.ExecutionClosure` (#829) * Parse duration field from flyteidl to `flytekit.models.execution.ExecutionClosure` Signed-off-by: Bernhard Stadlbauer * Add test for execution closure Signed-off-by: Bernhard Stadlbauer * Add tests to Flyte remote Signed-off-by: Bernhard Stadlbauer * Split execution test into with output and with error Signed-off-by: Bernhard Stadlbauer Co-authored-by: Bernhard Stadlbauer Signed-off-by: maximsmol --- flytekit/models/execution.py | 13 ++++- .../integration/remote/test_remote.py | 2 + tests/flytekit/unit/models/test_execution.py | 51 +++++++++++++++++++ 3 files changed, 65 insertions(+), 1 deletion(-) diff --git a/flytekit/models/execution.py b/flytekit/models/execution.py index a29bad80f1..8a46b09335 100644 --- a/flytekit/models/execution.py +++ b/flytekit/models/execution.py @@ -293,15 +293,17 @@ def from_flyte_idl(cls, pb): class ExecutionClosure(_common_models.FlyteIdlEntity): - def __init__(self, phase, started_at, error=None, outputs=None): + def __init__(self, phase, started_at, duration, error=None, outputs=None): """ :param int phase: From the flytekit.models.core.execution.WorkflowExecutionPhase enum :param datetime.datetime started_at: + :param datetime.timedelta duration: Duration for which the execution has been running. :param flytekit.models.core.execution.ExecutionError error: :param LiteralMapBlob outputs: """ self._phase = phase self._started_at = started_at + self._duration = duration self._error = error self._outputs = outputs @@ -327,6 +329,13 @@ def started_at(self): """ return self._started_at + @property + def duration(self): + """ + :rtype: datetime.timedelta + """ + return self._duration + @property def outputs(self): """ @@ -344,6 +353,7 @@ def to_flyte_idl(self): outputs=self.outputs.to_flyte_idl() if self.outputs is not None else None, ) obj.started_at.FromDatetime(self.started_at.astimezone(_pytz.UTC).replace(tzinfo=None)) + obj.duration.FromTimedelta(self.duration) return obj @classmethod @@ -363,6 +373,7 @@ def from_flyte_idl(cls, pb2_object): outputs=outputs, phase=pb2_object.phase, started_at=pb2_object.started_at.ToDatetime().replace(tzinfo=_pytz.UTC), + duration=pb2_object.duration.ToTimedelta(), ) diff --git a/tests/flytekit/integration/remote/test_remote.py b/tests/flytekit/integration/remote/test_remote.py index 7784d76d82..829137c928 100644 --- a/tests/flytekit/integration/remote/test_remote.py +++ b/tests/flytekit/integration/remote/test_remote.py @@ -145,6 +145,8 @@ def test_fetch_execute_workflow(flyteclient, flyte_workflows_register): flyte_workflow = remote.fetch_workflow(name="workflows.basic.hello_world.my_wf", version=f"v{VERSION}") execution = remote.execute(flyte_workflow, {}, wait=True) assert execution.outputs["o0"] == "hello world" + assert isinstance(execution.closure.duration, datetime.timedelta) + assert execution.closure.duration > datetime.timedelta(seconds=1) execution_to_terminate = remote.execute(flyte_workflow, {}) remote.terminate(execution_to_terminate, cause="just because") diff --git a/tests/flytekit/unit/models/test_execution.py b/tests/flytekit/unit/models/test_execution.py index c0fdf5ba2a..8c1ac94fce 100644 --- a/tests/flytekit/unit/models/test_execution.py +++ b/tests/flytekit/unit/models/test_execution.py @@ -1,4 +1,7 @@ +import datetime + import pytest +import pytz from flytekit.models import common as _common_models from flytekit.models import execution as _execution @@ -15,6 +18,54 @@ ) +def test_execution_closure_with_output(): + test_datetime = datetime.datetime(year=2022, month=1, day=1, tzinfo=pytz.UTC) + test_timedelta = datetime.timedelta(seconds=10) + test_outputs = _execution.LiteralMapBlob(values=_OUTPUT_MAP, uri="http://foo/") + + obj = _execution.ExecutionClosure( + phase=_core_exec.WorkflowExecutionPhase.SUCCEEDED, + started_at=test_datetime, + duration=test_timedelta, + outputs=test_outputs, + ) + assert obj.phase == _core_exec.WorkflowExecutionPhase.SUCCEEDED + assert obj.started_at == test_datetime + assert obj.duration == test_timedelta + assert obj.outputs == test_outputs + obj2 = _execution.ExecutionClosure.from_flyte_idl(obj.to_flyte_idl()) + assert obj2 == obj + assert obj2.phase == _core_exec.WorkflowExecutionPhase.SUCCEEDED + assert obj2.started_at == test_datetime + assert obj2.duration == test_timedelta + assert obj2.outputs == test_outputs + + +def test_execution_closure_with_error(): + test_datetime = datetime.datetime(year=2022, month=1, day=1, tzinfo=pytz.UTC) + test_timedelta = datetime.timedelta(seconds=10) + test_error = _core_exec.ExecutionError( + code="foo", message="bar", error_uri="http://foobar", kind=_core_exec.ExecutionError.ErrorKind.USER + ) + + obj = _execution.ExecutionClosure( + phase=_core_exec.WorkflowExecutionPhase.SUCCEEDED, + started_at=test_datetime, + duration=test_timedelta, + error=test_error, + ) + assert obj.phase == _core_exec.WorkflowExecutionPhase.SUCCEEDED + assert obj.started_at == test_datetime + assert obj.duration == test_timedelta + assert obj.error == test_error + obj2 = _execution.ExecutionClosure.from_flyte_idl(obj.to_flyte_idl()) + assert obj2 == obj + assert obj2.phase == _core_exec.WorkflowExecutionPhase.SUCCEEDED + assert obj2.started_at == test_datetime + assert obj2.duration == test_timedelta + assert obj2.error == test_error + + def test_execution_metadata(): obj = _execution.ExecutionMetadata(_execution.ExecutionMetadata.ExecutionMode.MANUAL, "tester", 1) assert obj.mode == _execution.ExecutionMetadata.ExecutionMode.MANUAL From 05548b0f3c9b49d275a65e0da08f9432192e64b0 Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Mon, 31 Jan 2022 15:19:02 -0800 Subject: [PATCH 074/128] Gate new Structured Dataset feature & remove old config objects (#831) Signed-off-by: Yee Hing Tong --- .github/workflows/pythonbuild.yml | 8 +- Dockerfile.py310 | 1 + Dockerfile.py37 | 1 + Dockerfile.py38 | 1 + Dockerfile.py39 | 1 + Makefile | 7 +- flytekit/__init__.py | 13 +- flytekit/configuration/platform.py | 5 - flytekit/configuration/sdk.py | 63 +------- flytekit/types/structured/__init__.py | 38 ++--- .../types/structured/structured_dataset.py | 13 +- plugins/Makefile | 2 +- .../flytekitplugins/spark/__init__.py | 7 +- .../flytekitplugins/spark/schema.py | 47 +----- .../flytekitplugins/spark/sd_transformers.py | 49 ++++++ plugins/flytekit-sqlalchemy/Dockerfile | 1 + .../mock_flyte_repo/workflows/Dockerfile | 1 + tests/flytekit_compatibility/__init__.py | 0 .../test_schema_types.py | 53 +++++++ .../test_schema_usage_copied.py | 149 ++++++++++++++++++ .../test_structured_dataset.py | 8 + 21 files changed, 324 insertions(+), 144 deletions(-) create mode 100644 plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py create mode 100644 tests/flytekit_compatibility/__init__.py create mode 100644 tests/flytekit_compatibility/test_schema_types.py create mode 100644 tests/flytekit_compatibility/test_schema_usage_copied.py create mode 100644 tests/flytekit_compatibility/test_structured_dataset.py diff --git a/.github/workflows/pythonbuild.yml b/.github/workflows/pythonbuild.yml index 2eeabd109d..b5a239a5df 100644 --- a/.github/workflows/pythonbuild.yml +++ b/.github/workflows/pythonbuild.yml @@ -38,17 +38,17 @@ jobs: run: | python -m pip install --upgrade pip==21.2.4 setuptools wheel make setup${{ matrix.spark-version-suffix }} - git clone https://github.com/maximsmol/flyteidl.git && cd flyteidl && git checkout maximsmol-union-types && pip install . && cd .. pip freeze - name: Test with coverage run: | - coverage run -m pytest tests/flytekit/unit + coverage run -m pytest tests/flytekit_compatibility + FLYTE_SDK_USE_STRUCTURED_DATASET=TRUE coverage run -m pytest tests/flytekit/unit - name: Integration Tests with coverage # https://github.com/actions/runner/issues/241#issuecomment-577360161 shell: 'script -q -e -c "bash {0}"' run: | python -m pip install awscli - coverage run --append -m pytest tests/flytekit/integration + FLYTE_SDK_USE_STRUCTURED_DATASET=TRUE coverage run --append -m pytest tests/flytekit/integration - name: Codecov uses: codecov/codecov-action@v1 with: @@ -109,7 +109,7 @@ jobs: - name: Test with coverage run: | cd plugins/${{ matrix.plugin-names }} - coverage run -m pytest tests + FLYTE_SDK_USE_STRUCTURED_DATASET=TRUE coverage run -m pytest tests lint: runs-on: ubuntu-latest diff --git a/Dockerfile.py310 b/Dockerfile.py310 index 1dc6b117b0..8fc7ae9e3d 100644 --- a/Dockerfile.py310 +++ b/Dockerfile.py310 @@ -14,3 +14,4 @@ RUN pip install -U flytekit==$VERSION WORKDIR /app ENV FLYTE_INTERNAL_IMAGE "$DOCKER_IMAGE" +ENV FLYTE_SDK_USE_STRUCTURED_DATASET TRUE diff --git a/Dockerfile.py37 b/Dockerfile.py37 index 4eb531b8de..ae59a7d608 100644 --- a/Dockerfile.py37 +++ b/Dockerfile.py37 @@ -14,3 +14,4 @@ RUN pip install -U flytekit==$VERSION WORKDIR /app ENV FLYTE_INTERNAL_IMAGE "$DOCKER_IMAGE" +ENV FLYTE_SDK_USE_STRUCTURED_DATASET TRUE diff --git a/Dockerfile.py38 b/Dockerfile.py38 index 934f167748..bdc9acec05 100644 --- a/Dockerfile.py38 +++ b/Dockerfile.py38 @@ -14,3 +14,4 @@ RUN pip install -U flytekit==$VERSION WORKDIR /app ENV FLYTE_INTERNAL_IMAGE "$DOCKER_IMAGE" +ENV FLYTE_SDK_USE_STRUCTURED_DATASET TRUE diff --git a/Dockerfile.py39 b/Dockerfile.py39 index 1a4d617aa0..6e972b2589 100644 --- a/Dockerfile.py39 +++ b/Dockerfile.py39 @@ -14,3 +14,4 @@ RUN pip install -U flytekit==$VERSION WORKDIR /app ENV FLYTE_INTERNAL_IMAGE "$DOCKER_IMAGE" +ENV FLYTE_SDK_USE_STRUCTURED_DATASET TRUE diff --git a/Makefile b/Makefile index b1088aece7..10f4e5b232 100644 --- a/Makefile +++ b/Makefile @@ -46,12 +46,11 @@ spellcheck: ## Runs a spellchecker over all code and documentation codespell -L "te,raison,fo" --skip="./docs/build,./.git" .PHONY: test -test: lint ## Run tests - pytest tests/flytekit/unit +test: lint unit_test .PHONY: unit_test unit_test: - pytest tests/flytekit/unit + FLYTE_SDK_USE_STRUCTURED_DATASET=TRUE pytest tests/flytekit/unit tests/flytekit_compatibility requirements-spark2.txt: export CUSTOM_COMPILE_COMMAND := make requirements-spark2.txt requirements-spark2.txt: requirements-spark2.in install-piptools @@ -79,7 +78,7 @@ requirements: requirements.txt dev-requirements.txt requirements-spark2.txt doc- # TODO: Change this in the future to be all of flytekit .PHONY: coverage coverage: - coverage run -m pytest tests/flytekit/unit/core flytekit/types + FLYTE_SDK_USE_STRUCTURED_DATASET=TRUE coverage run -m pytest tests/flytekit/unit/core flytekit/types coverage report -m --include="flytekit/core/*,flytekit/types/*" PLACEHOLDER := "__version__\ =\ \"0.0.0+develop\"" diff --git a/flytekit/__init__.py b/flytekit/__init__.py index 0c2db39db0..64c0a26530 100644 --- a/flytekit/__init__.py +++ b/flytekit/__init__.py @@ -160,6 +160,7 @@ else: from importlib.metadata import entry_points +from flytekit.configuration.sdk import USE_STRUCTURED_DATASET from flytekit.core.base_sql_task import SQLTask from flytekit.core.base_task import SecurityContext, TaskMetadata, kwtypes from flytekit.core.checkpointer import Checkpoint @@ -187,11 +188,13 @@ from flytekit.models.literals import Blob, BlobMetadata, Literal, Scalar from flytekit.models.types import LiteralType from flytekit.types import directory, file, schema -from flytekit.types.structured.structured_dataset import ( - StructuredDataset, - StructuredDatasetFormat, - StructuredDatasetType, -) + +if USE_STRUCTURED_DATASET.get(): + from flytekit.types.structured.structured_dataset import ( + StructuredDataset, + StructuredDatasetFormat, + StructuredDatasetType, + ) __version__ = "0.0.0+develop" diff --git a/flytekit/configuration/platform.py b/flytekit/configuration/platform.py index 5c4061fa4f..eecbeda162 100644 --- a/flytekit/configuration/platform.py +++ b/flytekit/configuration/platform.py @@ -1,5 +1,4 @@ from flytekit.configuration import common as _config_common -from flytekit.core import constants as _constants URL = _config_common.FlyteStringConfigurationEntry("platform", "url") @@ -14,10 +13,6 @@ INSECURE = _config_common.FlyteBoolConfigurationEntry("platform", "insecure", default=False) -CLOUD_PROVIDER = _config_common.FlyteStringConfigurationEntry( - "platform", "cloud_provider", default=_constants.CloudProvider.AWS -) - AUTH = _config_common.FlyteBoolConfigurationEntry("platform", "auth", default=False) """ This config setting should not normally be filled in. Whether or not an admin server requires authentication should be diff --git a/flytekit/configuration/sdk.py b/flytekit/configuration/sdk.py index 7e5c8872ce..a142b29546 100644 --- a/flytekit/configuration/sdk.py +++ b/flytekit/configuration/sdk.py @@ -8,20 +8,6 @@ and execution of entities. """ -EXECUTION_ENGINE = _config_common.FlyteStringConfigurationEntry("sdk", "execution_engine", default="flyte") -""" -This is a comma-delimited list of package strings, in order, for resolving execution behavior. - -TODO: Explain how this would be used to extend the SDK -""" - -TYPE_ENGINES = _config_common.FlyteStringListConfigurationEntry("sdk", "type_engines", default=[]) -""" -This is a comma-delimited list of package strings, in order, for resolving type behavior. - -TODO: Explain how this would be used to extend the SDK -""" - LOCAL_SANDBOX = _config_common.FlyteStringConfigurationEntry( "sdk", "local_sandbox", @@ -32,47 +18,6 @@ clean up data in these directories. """ -SDK_PYTHON_VENV = _config_common.FlyteStringListConfigurationEntry("sdk", "python_venv", default=[]) -""" -This is a list of commands/args which will be prefixed to the entrypoint command by SDK. -""" - -ROLE = _config_common.FlyteStringConfigurationEntry("sdk", "role") -""" -This is the role the SDK will use by default to execute workflows. For example, in AWS this should be an IAM role -string. -""" - -NAME_FORMAT = _config_common.FlyteStringConfigurationEntry("sdk", "name_format", default="{module}.{name}") -""" -This is a Python format string which the SDK will use to generate names for discovered entities. The default is -'{module}.{name}' which will result in strings like 'package.module.name'. Any template portion of the string can only -include 'module' or 'name'. So '{name}' is valid, but '{key}' is not. -""" - -TASK_NAME_FORMAT = _config_common.FlyteStringConfigurationEntry("sdk", "task_name_format", fallback=NAME_FORMAT) -""" -This is a Python format string which the SDK will use to generate names for tasks. Any template portion of the -string can only include 'module' or 'name'. So '{name}' is valid, but '{key}' is not. If not specified, -we fall back to the configuration for :py:attr:`flytekit.configuration.sdk.NAME_FORMAT` -""" - -WORKFLOW_NAME_FORMAT = _config_common.FlyteStringConfigurationEntry("sdk", "workflow_name_format", fallback=NAME_FORMAT) -""" -This is a Python format string which the SDK will use to generate names for workflows. Any template portion of the -string can only include 'module' or 'name'. So '{name}' is valid, but '{key}' is not. If not specified, -we fall back to the configuration for :py:attr:`flytekit.configuration.sdk.NAME_FORMAT` -""" - -LAUNCH_PLAN_NAME_FORMAT = _config_common.FlyteStringConfigurationEntry( - "sdk", "launch_plan_name_format", fallback=NAME_FORMAT -) -""" -This is a Python format string which the SDK will use to generate names for launch plans. Any template portion of the -string can only include 'module' or 'name'. So '{name}' is valid, but '{key}' is not. If not specified, -we fall back to the configuration for :py:attr:`flytekit.configuration.sdk.NAME_FORMAT` -""" - LOGGING_LEVEL = _config_common.FlyteIntegerConfigurationEntry("sdk", "logging_level", default=20) """ This is the default logging level for the Python logging library and will be set before user code runs. @@ -85,9 +30,9 @@ This is the parquet engine to use when reading data from parquet files. """ -FAST_REGISTRATION_DIR = _config_common.FlyteStringConfigurationEntry("sdk", "fast_registration_dir") +# Feature Gate +USE_STRUCTURED_DATASET = _config_common.FlyteBoolConfigurationEntry("sdk", "use_structured_dataset", default=False) """ -This is the remote directory where fast-registered code will be uploaded to. -Users calling fast-execute need write permission to this directory. -Furthermore, it is important that whichever role executes your workflow has read access to this directory. +Note: This gate will be switched to True at some point in the future. Definitely by 1.0, if not v0.31.0. + """ diff --git a/flytekit/types/structured/__init__.py b/flytekit/types/structured/__init__.py index da80c89016..c4fd015c57 100644 --- a/flytekit/types/structured/__init__.py +++ b/flytekit/types/structured/__init__.py @@ -1,21 +1,23 @@ +from flytekit.configuration.sdk import USE_STRUCTURED_DATASET from flytekit.loggers import logger -from .basic_dfs import ( - ArrowToParquetEncodingHandler, - PandasToParquetEncodingHandler, - ParquetToArrowDecodingHandler, - ParquetToPandasDecodingHandler, -) - -try: - from .bigquery import ( - ArrowToBQEncodingHandlers, - BQToArrowDecodingHandler, - BQToPandasDecodingHandler, - PandasToBQEncodingHandlers, - ) -except ImportError: - logger.info( - "We won't register bigquery handler for structured dataset because " - "we can't find the packages google-cloud-bigquery-storage and google-cloud-bigquery" +if USE_STRUCTURED_DATASET.get(): + from .basic_dfs import ( + ArrowToParquetEncodingHandler, + PandasToParquetEncodingHandler, + ParquetToArrowDecodingHandler, + ParquetToPandasDecodingHandler, ) + + try: + from .bigquery import ( + ArrowToBQEncodingHandlers, + BQToArrowDecodingHandler, + BQToPandasDecodingHandler, + PandasToBQEncodingHandlers, + ) + except ImportError: + logger.info( + "We won't register bigquery handler for structured dataset because " + "we can't find the packages google-cloud-bigquery-storage and google-cloud-bigquery" + ) diff --git a/flytekit/types/structured/structured_dataset.py b/flytekit/types/structured/structured_dataset.py index a36e50b976..ffefd101dd 100644 --- a/flytekit/types/structured/structured_dataset.py +++ b/flytekit/types/structured/structured_dataset.py @@ -22,6 +22,7 @@ import numpy as _np import pyarrow as pa +from flytekit.configuration.sdk import USE_STRUCTURED_DATASET from flytekit.core.context_manager import FlyteContext, FlyteContextManager from flytekit.core.type_engine import TypeTransformer from flytekit.extend import TypeEngine @@ -381,6 +382,10 @@ def register_handler(self, h: Handlers, default_for_type: Optional[bool] = True, The string "://" should not be present in any handler's protocol so we don't check for it. """ + if not USE_STRUCTURED_DATASET.get(): + logger.info(f"Structured datasets not enabled, not registering handler {h}") + return + lowest_level = self._handler_finder(h) if h.supported_format in lowest_level and override is False: raise ValueError(f"Already registered a handler for {(h.python_type, h.protocol, h.supported_format)}") @@ -717,5 +722,9 @@ def guess_python_type(self, literal_type: LiteralType) -> Type[T]: raise ValueError(f"StructuredDatasetTransformerEngine cannot reverse {literal_type}") -FLYTE_DATASET_TRANSFORMER = StructuredDatasetTransformerEngine() -TypeEngine.register(FLYTE_DATASET_TRANSFORMER) +if USE_STRUCTURED_DATASET.get(): + logger.debug("Structured dataset module load... using structured datasets!") + FLYTE_DATASET_TRANSFORMER = StructuredDatasetTransformerEngine() + TypeEngine.register(FLYTE_DATASET_TRANSFORMER) +else: + logger.debug("Structured dataset module load... not using structured datasets") diff --git a/plugins/Makefile b/plugins/Makefile index 9131836c77..94dabbde85 100644 --- a/plugins/Makefile +++ b/plugins/Makefile @@ -1,6 +1,6 @@ .PHONY: test test: - find . -maxdepth 1 -type d | grep 'flytekit-' | xargs -L1 pytest + FLYTE_SDK_USE_STRUCTURED_DATASET=TRUE find . -maxdepth 1 -type d | grep 'flytekit-' | xargs -L1 pytest .PHONY: build_all_plugins build_all_plugins: diff --git a/plugins/flytekit-spark/flytekitplugins/spark/__init__.py b/plugins/flytekit-spark/flytekitplugins/spark/__init__.py index 145497b030..d239632248 100644 --- a/plugins/flytekit-spark/flytekitplugins/spark/__init__.py +++ b/plugins/flytekit-spark/flytekitplugins/spark/__init__.py @@ -1,2 +1,7 @@ -from .schema import ParquetToSparkDecodingHandler, SparkToParquetEncodingHandler +from flytekit.configuration.sdk import USE_STRUCTURED_DATASET + +from .schema import SparkDataFrameSchemaReader, SparkDataFrameSchemaWriter, SparkDataFrameTransformer # noqa from .task import Spark, new_spark_session + +if USE_STRUCTURED_DATASET.get(): + from .sd_transformers import ParquetToSparkDecodingHandler, SparkToParquetEncodingHandler diff --git a/plugins/flytekit-spark/flytekitplugins/spark/schema.py b/plugins/flytekit-spark/flytekitplugins/spark/schema.py index 72dad4fc92..1cae101295 100644 --- a/plugins/flytekit-spark/flytekitplugins/spark/schema.py +++ b/plugins/flytekit-spark/flytekitplugins/spark/schema.py @@ -2,21 +2,12 @@ from typing import Type import pyspark -from pyspark.sql.dataframe import DataFrame from flytekit import FlyteContext from flytekit.extend import T, TypeEngine, TypeTransformer -from flytekit.models import literals -from flytekit.models.literals import Literal, Scalar, Schema, StructuredDatasetMetadata -from flytekit.models.types import LiteralType, SchemaType, StructuredDatasetType +from flytekit.models.literals import Literal, Scalar, Schema +from flytekit.models.types import LiteralType, SchemaType from flytekit.types.schema import SchemaEngine, SchemaFormat, SchemaHandler, SchemaReader, SchemaWriter -from flytekit.types.structured.structured_dataset import ( - FLYTE_DATASET_TRANSFORMER, - PARQUET, - StructuredDataset, - StructuredDatasetDecoder, - StructuredDatasetEncoder, -) class SparkDataFrameSchemaReader(SchemaReader[pyspark.sql.DataFrame]): @@ -106,37 +97,3 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: # %% # This makes pyspark.DataFrame as a supported output/input type with flytekit. TypeEngine.register(SparkDataFrameTransformer()) - - -class SparkToParquetEncodingHandler(StructuredDatasetEncoder): - def __init__(self, protocol: str): - super().__init__(DataFrame, protocol, PARQUET) - - def encode( - self, - ctx: FlyteContext, - structured_dataset: StructuredDataset, - structured_dataset_type: StructuredDatasetType, - ) -> literals.StructuredDataset: - path = typing.cast(str, structured_dataset.uri) or ctx.file_access.get_random_remote_directory() - df = typing.cast(DataFrame, structured_dataset.dataframe) - df.write.mode("overwrite").parquet(path) - return literals.StructuredDataset(uri=path, metadata=StructuredDatasetMetadata(structured_dataset_type)) - - -class ParquetToSparkDecodingHandler(StructuredDatasetDecoder): - def __init__(self, protocol: str): - super().__init__(DataFrame, protocol, PARQUET) - - def decode( - self, - ctx: FlyteContext, - flyte_value: literals.StructuredDataset, - ) -> DataFrame: - user_ctx = FlyteContext.current_context().user_space_params - return user_ctx.spark_session.read.parquet(flyte_value.uri) - - -for protocol in ["/", "s3"]: - FLYTE_DATASET_TRANSFORMER.register_handler(SparkToParquetEncodingHandler(protocol), default_for_type=True) - FLYTE_DATASET_TRANSFORMER.register_handler(ParquetToSparkDecodingHandler(protocol), default_for_type=True) diff --git a/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py b/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py new file mode 100644 index 0000000000..e0b1c7b41e --- /dev/null +++ b/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py @@ -0,0 +1,49 @@ +import typing + +from pyspark.sql.dataframe import DataFrame + +from flytekit import FlyteContext +from flytekit.models import literals +from flytekit.models.literals import StructuredDatasetMetadata +from flytekit.models.types import StructuredDatasetType +from flytekit.types.structured.structured_dataset import ( + FLYTE_DATASET_TRANSFORMER, + PARQUET, + StructuredDataset, + StructuredDatasetDecoder, + StructuredDatasetEncoder, +) + + +class SparkToParquetEncodingHandler(StructuredDatasetEncoder): + def __init__(self, protocol: str): + super().__init__(DataFrame, protocol, PARQUET) + + def encode( + self, + ctx: FlyteContext, + structured_dataset: StructuredDataset, + structured_dataset_type: StructuredDatasetType, + ) -> literals.StructuredDataset: + path = typing.cast(str, structured_dataset.uri) or ctx.file_access.get_random_remote_directory() + df = typing.cast(DataFrame, structured_dataset.dataframe) + df.write.mode("overwrite").parquet(path) + return literals.StructuredDataset(uri=path, metadata=StructuredDatasetMetadata(structured_dataset_type)) + + +class ParquetToSparkDecodingHandler(StructuredDatasetDecoder): + def __init__(self, protocol: str): + super().__init__(DataFrame, protocol, PARQUET) + + def decode( + self, + ctx: FlyteContext, + flyte_value: literals.StructuredDataset, + ) -> DataFrame: + user_ctx = FlyteContext.current_context().user_space_params + return user_ctx.spark_session.read.parquet(flyte_value.uri) + + +for protocol in ["/", "s3"]: + FLYTE_DATASET_TRANSFORMER.register_handler(SparkToParquetEncodingHandler(protocol), default_for_type=True) + FLYTE_DATASET_TRANSFORMER.register_handler(ParquetToSparkDecodingHandler(protocol), default_for_type=True) diff --git a/plugins/flytekit-sqlalchemy/Dockerfile b/plugins/flytekit-sqlalchemy/Dockerfile index b1c47ffd1e..74c9678746 100644 --- a/plugins/flytekit-sqlalchemy/Dockerfile +++ b/plugins/flytekit-sqlalchemy/Dockerfile @@ -5,6 +5,7 @@ ENV VENV /opt/venv ENV LANG C.UTF-8 ENV LC_ALL C.UTF-8 ENV PYTHONPATH /app +ENV FLYTE_SDK_USE_STRUCTURED_DATASET TRUE RUN pip install awscli RUN pip install gsutil diff --git a/tests/flytekit/integration/remote/mock_flyte_repo/workflows/Dockerfile b/tests/flytekit/integration/remote/mock_flyte_repo/workflows/Dockerfile index 7e5d01829f..5548412616 100644 --- a/tests/flytekit/integration/remote/mock_flyte_repo/workflows/Dockerfile +++ b/tests/flytekit/integration/remote/mock_flyte_repo/workflows/Dockerfile @@ -6,6 +6,7 @@ ENV VENV /opt/venv ENV LANG C.UTF-8 ENV LC_ALL C.UTF-8 ENV PYTHONPATH /root +ENV FLYTE_SDK_USE_STRUCTURED_DATASET TRUE # This is necessary for opencv to work RUN apt-get update && apt-get install -y libsm6 libxext6 libxrender-dev ffmpeg build-essential diff --git a/tests/flytekit_compatibility/__init__.py b/tests/flytekit_compatibility/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/flytekit_compatibility/test_schema_types.py b/tests/flytekit_compatibility/test_schema_types.py new file mode 100644 index 0000000000..6711ec34de --- /dev/null +++ b/tests/flytekit_compatibility/test_schema_types.py @@ -0,0 +1,53 @@ +from datetime import datetime, timedelta + +import pytest + +from flytekit import kwtypes +from flytekit.core import context_manager +from flytekit.core.context_manager import ExecutionState +from flytekit.core.type_engine import TypeEngine +from flytekit.types.schema import FlyteSchema, SchemaFormat +from flytekit.types.schema.types import FlyteSchemaTransformer + + +def test_typed_schema(): + s = FlyteSchema[kwtypes(x=int, y=float)] + assert s.format() == SchemaFormat.PARQUET + assert s.columns() == {"x": int, "y": float} + + +def test_assert_type(): + ctx = context_manager.FlyteContextManager.current_context() + with context_manager.FlyteContextManager.with_context( + ctx.with_execution_state(ctx.new_execution_state().with_params(mode=ExecutionState.Mode.TASK_EXECUTION)) + ) as ctx: + schema = FlyteSchema[kwtypes(x=int, y=float)] + fst = FlyteSchemaTransformer() + lt = fst.get_literal_type(schema) + with pytest.raises(ValueError, match="DataFrames of type are not supported currently"): + TypeEngine.to_literal(ctx, 3, schema, lt) + + +def test_schema_back_and_forth(): + orig = FlyteSchema[kwtypes(TrackId=int, Name=str)] + lt = TypeEngine.to_literal_type(orig) + pt = TypeEngine.guess_python_type(lt) + lt2 = TypeEngine.to_literal_type(pt) + assert lt == lt2 + + +def test_remaining_prims(): + orig = FlyteSchema[kwtypes(my_dt=datetime, my_td=timedelta, my_b=bool)] + lt = TypeEngine.to_literal_type(orig) + pt = TypeEngine.guess_python_type(lt) + lt2 = TypeEngine.to_literal_type(pt) + assert lt == lt2 + + +def test_bad_conversion(): + orig = FlyteSchema[kwtypes(my_custom=bool)] + lt = TypeEngine.to_literal_type(orig) + # Make a not real column type + lt.schema.columns[0]._type = 15 + with pytest.raises(ValueError): + TypeEngine.guess_python_type(lt) diff --git a/tests/flytekit_compatibility/test_schema_usage_copied.py b/tests/flytekit_compatibility/test_schema_usage_copied.py new file mode 100644 index 0000000000..dfeee49cf3 --- /dev/null +++ b/tests/flytekit_compatibility/test_schema_usage_copied.py @@ -0,0 +1,149 @@ +import datetime +from dataclasses import dataclass + +import pandas as pd +from dataclasses_json import dataclass_json + +from flytekit import SQLTask, kwtypes +from flytekit.core import context_manager +from flytekit.core.task import TaskMetadata, task +from flytekit.core.testing import patch, task_mock +from flytekit.core.workflow import workflow +from flytekit.types.schema import FlyteSchema, SchemaOpenMode + + +def test_wf1_with_sql(): + sql = SQLTask( + "my-query", + query_template="SELECT * FROM hive.city.fact_airport_sessions WHERE ds = '{{ .Inputs.ds }}' LIMIT 10", + inputs=kwtypes(ds=datetime.datetime), + outputs=kwtypes(results=FlyteSchema), + metadata=TaskMetadata(retries=2), + ) + + @task + def t1() -> datetime.datetime: + return datetime.datetime.now() + + @workflow + def my_wf() -> FlyteSchema: + dt = t1() + return sql(ds=dt) + + with task_mock(sql) as mock: + mock.return_value = pd.DataFrame(data={"x": [1, 2], "y": ["3", "4"]}) + assert (my_wf().open().all() == pd.DataFrame(data={"x": [1, 2], "y": ["3", "4"]})).all().all() + assert context_manager.FlyteContextManager.size() == 1 + + +def test_wf1_with_sql_with_patch(): + sql = SQLTask( + "my-query", + query_template="SELECT * FROM hive.city.fact_airport_sessions WHERE ds = '{{ .Inputs.ds }}' LIMIT 10", + inputs=kwtypes(ds=datetime.datetime), + outputs=kwtypes(results=FlyteSchema), + metadata=TaskMetadata(retries=2), + ) + + @task + def t1() -> datetime.datetime: + return datetime.datetime.now() + + @workflow + def my_wf() -> FlyteSchema: + dt = t1() + return sql(ds=dt) + + @patch(sql) + def test_user_demo_test(mock_sql): + mock_sql.return_value = pd.DataFrame(data={"x": [1, 2], "y": ["3", "4"]}) + assert (my_wf().open().all() == pd.DataFrame(data={"x": [1, 2], "y": ["3", "4"]})).all().all() + + # Have to call because tests inside tests don't run + test_user_demo_test() + assert context_manager.FlyteContextManager.size() == 1 + + +def test_wf_typed_schema(): + schema1 = FlyteSchema[kwtypes(x=int, y=str)] + + @task + def t1() -> schema1: + s = schema1() + s.open().write(pd.DataFrame(data={"x": [1, 2], "y": ["3", "4"]})) + return s + + @task + def t2(s: FlyteSchema[kwtypes(x=int, y=str)]) -> FlyteSchema[kwtypes(x=int)]: + df = s.open().all() + return df[s.column_names()[:-1]] + + @workflow + def wf() -> FlyteSchema[kwtypes(x=int)]: + return t2(s=t1()) + + w = t1() + assert w is not None + df = w.open(override_mode=SchemaOpenMode.READ).all() + result_df = df.reset_index(drop=True) == pd.DataFrame(data={"x": [1, 2], "y": ["3", "4"]}).reset_index(drop=True) + assert result_df.all().all() + + df = t2(s=w.as_readonly()) + df = df.open(override_mode=SchemaOpenMode.READ).all() + result_df = df.reset_index(drop=True) == pd.DataFrame(data={"x": [1, 2]}).reset_index(drop=True) + assert result_df.all().all() + + x = wf() + df = x.open().all() + result_df = df.reset_index(drop=True) == pd.DataFrame(data={"x": [1, 2]}).reset_index(drop=True) + assert result_df.all().all() + + +def test_wf_schema_to_df(): + schema1 = FlyteSchema[kwtypes(x=int, y=str)] + + @task + def t1() -> schema1: + s = schema1() + s.open().write(pd.DataFrame(data={"x": [1, 2], "y": ["3", "4"]})) + return s + + @task + def t2(df: pd.DataFrame) -> int: + return len(df.columns.values) + + @workflow + def wf() -> int: + return t2(df=t1()) + + x = wf() + assert x == 2 + + +def test_flyte_schema_dataclass(): + TestSchema = FlyteSchema[kwtypes(some_str=str)] + + @dataclass_json + @dataclass + class InnerResult: + number: int + schema: TestSchema + + @dataclass_json + @dataclass + class Result: + result: InnerResult + schema: TestSchema + + schema = TestSchema() + + @task + def t1(x: int) -> Result: + + return Result(result=InnerResult(number=x, schema=schema), schema=schema) + + @workflow + def wf(x: int) -> Result: + return t1(x=x) + + assert wf(x=10) == Result(result=InnerResult(number=10, schema=schema), schema=schema) diff --git a/tests/flytekit_compatibility/test_structured_dataset.py b/tests/flytekit_compatibility/test_structured_dataset.py new file mode 100644 index 0000000000..20965cb802 --- /dev/null +++ b/tests/flytekit_compatibility/test_structured_dataset.py @@ -0,0 +1,8 @@ +import pandas as pd + +from flytekit.core.type_engine import TypeEngine + + +def test_pandas_is_schema_with_flag(): + lt = TypeEngine.to_literal_type(pd.DataFrame) + assert lt.schema is not None From 71b665d99c2c5ef661f08ab9175ae7e4f431b65d Mon Sep 17 00:00:00 2001 From: Ketan Umare <16888709+kumare3@users.noreply.github.com> Date: Wed, 2 Feb 2022 17:47:10 -0800 Subject: [PATCH 075/128] Fixing out of order for conditional outputs (#843) Signed-off-by: maximsmol --- flytekit/core/condition.py | 34 ++++++++++++++++++++++------------ 1 file changed, 22 insertions(+), 12 deletions(-) diff --git a/flytekit/core/condition.py b/flytekit/core/condition.py index a38d0e2ab1..9b2670a683 100644 --- a/flytekit/core/condition.py +++ b/flytekit/core/condition.py @@ -114,29 +114,39 @@ def end_branch(self) -> Optional[Union[Condition, Promise, Tuple[Promise], VoidP def if_(self, expr: bool) -> Case: return self._condition._if(expr) - def compute_output_set(self) -> typing.Optional[typing.Set[str]]: + def compute_output_vars(self) -> typing.Optional[typing.List[str]]: """ Computes and returns the minimum set of outputs for this conditional block, based on all the cases that have been registered """ - output_var_sets: typing.List[typing.Set[str]] = [] + output_vars: typing.List[str] = [] + output_vars_set = set() for c in self._cases: if c.output_promise is None and c.err is None: # One node returns a void output and no error, we will default to None return return None if c.output_promise is not None: + var = [] if isinstance(c.output_promise, tuple): - output_var_sets.append(set([i.var for i in c.output_promise])) + var = [i.var for i in c.output_promise] else: - output_var_sets.append({c.output_promise.var}) - curr = output_var_sets[0] - if len(output_var_sets) > 1: - for x in output_var_sets[1:]: - curr = curr.intersection(x) - return curr + var = [c.output_promise.var] + curr_set = set(var) + if not output_vars: + output_vars = var + output_vars_set = curr_set + else: + output_vars_set = output_vars_set.intersection(curr_set) + new_output_var = [] + for v in output_vars: + if v in output_vars_set: + new_output_var.append(v) + output_vars = new_output_var + + return output_vars def _compute_outputs(self, n: Node) -> Optional[Union[Promise, Tuple[Promise], VoidPromise]]: - curr = self.compute_output_set() + curr = self.compute_output_vars() if curr is None: return VoidPromise(n.id) promises = [Promise(var=x, val=NodeOutput(node=n, var=x)) for x in curr] @@ -197,7 +207,7 @@ def _compute_outputs(self, selected_output_promise) -> Optional[Union[Tuple[Prom """ For the local execution case only returns the least common set of outputs """ - curr = self.compute_output_set() + curr = self.compute_output_vars() if curr is None: return VoidPromise(self.name) if not isinstance(selected_output_promise, tuple): @@ -221,7 +231,7 @@ def end_branch(self) -> Optional[Union[Condition, Tuple[Promise], Promise, VoidP """ if self._last_case: FlyteContextManager.pop_context() - curr = self.compute_output_set() + curr = self.compute_output_vars() if curr is None: return VoidPromise(self.name) promises = [Promise(var=x, val=None) for x in curr] From 3db68240108806151229d8af5e057496a1f200ce Mon Sep 17 00:00:00 2001 From: Eduardo Apolinario <653394+eapolinario@users.noreply.github.com> Date: Mon, 7 Feb 2022 16:38:45 -0800 Subject: [PATCH 076/128] Set default values to map task template (#841) * Set sane defaults in map task templates Signed-off-by: Eduardo Apolinario * Remove unused method Signed-off-by: Eduardo Apolinario * Put ArrayJob.from_dict back Signed-off-by: Eduardo Apolinario * Define parallelism=0 as unbounded Signed-off-by: Eduardo Apolinario * Remove special case to handle 0 Signed-off-by: Eduardo Apolinario Co-authored-by: Eduardo Apolinario Signed-off-by: maximsmol --- flytekit/core/map_task.py | 4 +- flytekit/models/array_job.py | 31 ++++++++++---- tests/flytekit/unit/core/test_map_task.py | 50 +++++++++++++++-------- 3 files changed, 57 insertions(+), 28 deletions(-) diff --git a/flytekit/core/map_task.py b/flytekit/core/map_task.py index f760be5d3c..4ec9f64a2d 100644 --- a/flytekit/core/map_task.py +++ b/flytekit/core/map_task.py @@ -209,7 +209,7 @@ def _raw_execute(self, **kwargs) -> Any: return outputs -def map_task(task_function: PythonFunctionTask, concurrency: int = None, min_success_ratio: float = None, **kwargs): +def map_task(task_function: PythonFunctionTask, concurrency: int = 0, min_success_ratio: float = 1.0, **kwargs): """ Use a map task for parallelizable tasks that run across a list of an input type. A map task can be composed of any individual :py:class:`flytekit.PythonFunctionTask`. @@ -231,7 +231,7 @@ def map_task(task_function: PythonFunctionTask, concurrency: int = None, min_suc :param task_function: This argument is implicitly passed and represents the repeatable function :param concurrency: If specified, this limits the number of mapped tasks than can run in parallel to the given batch size. If the size of the input exceeds the concurrency value, then multiple batches will be run serially until - all inputs are processed. + all inputs are processed. If left unspecified, this means unbounded concurrency. :param min_success_ratio: If specified, this determines the minimum fraction of total jobs which can complete successfully before terminating this task and marking it successful. diff --git a/flytekit/models/array_job.py b/flytekit/models/array_job.py index 4e4bf99cc7..2c86acdd7e 100644 --- a/flytekit/models/array_job.py +++ b/flytekit/models/array_job.py @@ -70,13 +70,21 @@ def to_dict(self): """ :rtype: dict[T, Text] """ - return _json_format.MessageToDict( - _array_job.ArrayJob( + array_job = None + if self.min_successes is not None: + array_job = _array_job.ArrayJob( parallelism=self.parallelism, size=self.size, min_successes=self.min_successes, ) - ) + elif self.min_success_ratio is not None: + array_job = _array_job.ArrayJob( + parallelism=self.parallelism, + size=self.size, + min_success_ratio=self.min_success_ratio, + ) + + return _json_format.MessageToDict(array_job) @classmethod def from_dict(cls, idl_dict): @@ -86,8 +94,15 @@ def from_dict(cls, idl_dict): """ pb2_object = _json_format.Parse(_json.dumps(idl_dict), _array_job.ArrayJob()) - return cls( - parallelism=pb2_object.parallelism, - size=pb2_object.size, - min_successes=pb2_object.min_successes, - ) + if pb2_object.HasField("min_successes"): + return cls( + parallelism=pb2_object.parallelism, + size=pb2_object.size, + min_successes=pb2_object.min_successes, + ) + else: + return cls( + parallelism=pb2_object.parallelism, + size=pb2_object.size, + min_success_ratio=pb2_object.min_success_ratio, + ) diff --git a/tests/flytekit/unit/core/test_map_task.py b/tests/flytekit/unit/core/test_map_task.py index d1f95852c1..4eb44d6e76 100644 --- a/tests/flytekit/unit/core/test_map_task.py +++ b/tests/flytekit/unit/core/test_map_task.py @@ -12,6 +12,18 @@ from flytekit.tools.translator import get_serializable +@pytest.fixture +def serialization_settings(): + default_img = Image(name="default", fqn="test", tag="tag") + return context_manager.SerializationSettings( + project="project", + domain="domain", + version="version", + env=None, + image_config=ImageConfig(default_image=default_img, images=[default_img]), + ) + + @task def t1(a: int) -> str: b = a + 2 @@ -54,18 +66,12 @@ def test_map_task_types(): _ = map_task(t1, metadata=TaskMetadata(retries=1))(a=["invalid", "args"]) -def test_serialization(): +def test_serialization(serialization_settings): maptask = map_task(t1, metadata=TaskMetadata(retries=1)) - default_img = Image(name="default", fqn="test", tag="tag") - serialization_settings = context_manager.SerializationSettings( - project="project", - domain="domain", - version="version", - env=None, - image_config=ImageConfig(default_image=default_img, images=[default_img]), - ) task_spec = get_serializable(OrderedDict(), serialization_settings, maptask) + # By default all map_task tasks will have their custom fields set. + assert task_spec.template.custom["minSuccessRatio"] == 1.0 assert task_spec.template.type == "container_array" assert task_spec.template.task_type_version == 1 assert task_spec.template.container.args == [ @@ -90,7 +96,23 @@ def test_serialization(): ] -def test_serialization_workflow_def(): +@pytest.mark.parametrize( + "custom_fields_dict, expected_custom_fields", + [ + ({}, {"minSuccessRatio": 1.0}), + ({"concurrency": 99}, {"parallelism": "99", "minSuccessRatio": 1.0}), + ({"min_success_ratio": 0.271828}, {"minSuccessRatio": 0.271828}), + ({"concurrency": 42, "min_success_ratio": 0.31415}, {"parallelism": "42", "minSuccessRatio": 0.31415}), + ], +) +def test_serialization_of_custom_fields(custom_fields_dict, expected_custom_fields, serialization_settings): + maptask = map_task(t1, **custom_fields_dict) + task_spec = get_serializable(OrderedDict(), serialization_settings, maptask) + + assert task_spec.template.custom == expected_custom_fields + + +def test_serialization_workflow_def(serialization_settings): @task def complex_task(a: int) -> str: b = a + 2 @@ -106,14 +128,6 @@ def w1(a: typing.List[int]) -> typing.List[str]: def w2(a: typing.List[int]) -> typing.List[str]: return map_task(complex_task, metadata=TaskMetadata(retries=2))(a=a) - default_img = Image(name="default", fqn="test", tag="tag") - serialization_settings = context_manager.SerializationSettings( - project="project", - domain="domain", - version="version", - env=None, - image_config=ImageConfig(default_image=default_img, images=[default_img]), - ) serialized_control_plane_entities = OrderedDict() wf1_spec = get_serializable(serialized_control_plane_entities, serialization_settings, w1) assert wf1_spec.template is not None From 110c9007f6471ad300445c512c9a7ec54db7268d Mon Sep 17 00:00:00 2001 From: Eduardo Apolinario <653394+eapolinario@users.noreply.github.com> Date: Tue, 8 Feb 2022 09:15:49 -0800 Subject: [PATCH 077/128] Regenerate requirements files - dependencies dropping support for python 3.7 (#838) * Regenerate requirements files Signed-off-by: Eduardo Apolinario * Put restrictions on numpy and pandas versions Signed-off-by: Eduardo Apolinario * Use --use-deprecated=legacy-resolver Signed-off-by: Eduardo Apolinario * use pip==22.0.3 everywhere Signed-off-by: Eduardo Apolinario * Remove --use-deprecated=legacy-resolver Signed-off-by: Eduardo Apolinario * Relax click Signed-off-by: Eduardo Apolinario * Regenerated plugins requirements Signed-off-by: Eduardo Apolinario Co-authored-by: Eduardo Apolinario --- .github/workflows/pythonbuild.yml | 2 - Makefile | 2 +- dev-requirements.txt | 109 ++++--- doc-requirements.txt | 269 +++--------------- plugins/flytekit-aws-athena/requirements.txt | 60 ++-- .../flytekit-aws-sagemaker/requirements.txt | 85 +++--- plugins/flytekit-bigquery/requirements.txt | 87 +++--- plugins/flytekit-data-fsspec/requirements.txt | 62 ++-- plugins/flytekit-dolt/requirements.txt | 62 ++-- .../requirements.txt | 128 +++++---- plugins/flytekit-hive/requirements.txt | 60 ++-- plugins/flytekit-k8s-pod/requirements.txt | 74 ++--- plugins/flytekit-kf-mpi/requirements.txt | 60 ++-- plugins/flytekit-kf-pytorch/requirements.txt | 60 ++-- .../flytekit-kf-tensorflow/requirements.txt | 60 ++-- plugins/flytekit-modin/requirements.txt | 94 +++--- plugins/flytekit-pandera/requirements.txt | 66 ++--- plugins/flytekit-papermill/requirements.txt | 120 ++++---- plugins/flytekit-snowflake/requirements.txt | 60 ++-- plugins/flytekit-spark/requirements.txt | 64 ++--- plugins/flytekit-sqlalchemy/requirements.txt | 62 ++-- requirements-spark2.txt | 250 ++-------------- requirements.in | 4 + requirements.txt | 250 ++-------------- setup.py | 2 +- .../workflows/requirements.txt | 54 ++-- 26 files changed, 836 insertions(+), 1370 deletions(-) diff --git a/.github/workflows/pythonbuild.yml b/.github/workflows/pythonbuild.yml index b5a239a5df..eda870cff4 100644 --- a/.github/workflows/pythonbuild.yml +++ b/.github/workflows/pythonbuild.yml @@ -36,7 +36,6 @@ jobs: key: ${{ format('{0}-pip-{1}', runner.os, hashFiles('dev-requirements.txt', format('requirements{0}.txt', matrix.spark-version-suffix))) }} - name: Install dependencies run: | - python -m pip install --upgrade pip==21.2.4 setuptools wheel make setup${{ matrix.spark-version-suffix }} pip freeze - name: Test with coverage @@ -99,7 +98,6 @@ jobs: key: ${{ format('{0}-pip-{1}', runner.os, hashFiles('dev-requirements.txt', format('plugins/{0}/requirements.txt', matrix.plugin-names ))) }} - name: Install dependencies run: | - python -m pip install --upgrade pip==21.2.4 setuptools wheel make setup cd plugins/${{ matrix.plugin-names }} pip install -e . diff --git a/Makefile b/Makefile index 10f4e5b232..b5464e205e 100644 --- a/Makefile +++ b/Makefile @@ -12,7 +12,7 @@ help: .PHONY: install-piptools install-piptools: - pip install -U pip-tools pip==21.2.4 + pip install -U pip-tools setuptools wheel pip==22.0.3 .PHONY: update_boilerplate update_boilerplate: diff --git a/dev-requirements.txt b/dev-requirements.txt index 7c88c15106..28c22afa15 100644 --- a/dev-requirements.txt +++ b/dev-requirements.txt @@ -2,13 +2,13 @@ # This file is autogenerated by pip-compile with python 3.9 # To update, run: # -# pip-compile dev-requirements.in +# make dev-requirements.txt # -e file:.#egg=flytekit # via # -c requirements.txt # pytest-flyte -arrow==1.2.1 +arrow==1.2.2 # via # -c requirements.txt # jinja2-time @@ -18,17 +18,13 @@ attrs==20.3.0 # jsonschema # pytest # pytest-docker -backports.entry-points-selectable==1.1.1 - # via virtualenv bcrypt==3.2.0 - # via - # -c requirements.txt - # paramiko + # via paramiko binaryornot==0.4.4 # via # -c requirements.txt # cookiecutter -cachetools==4.2.4 +cachetools==5.0.0 # via google-auth certifi==2021.10.8 # via @@ -46,7 +42,7 @@ chardet==4.0.0 # via # -c requirements.txt # binaryornot -charset-normalizer==2.0.9 +charset-normalizer==2.0.11 # via # -c requirements.txt # requests @@ -54,7 +50,7 @@ checksumdir==1.2.0 # via # -c requirements.txt # flytekit -click==7.1.2 +click==8.0.3 # via # -c requirements.txt # cookiecutter @@ -69,9 +65,9 @@ cookiecutter==1.7.3 # via # -c requirements.txt # flytekit -coverage[toml]==6.2 +coverage[toml]==6.3.1 # via -r dev-requirements.in -croniter==1.1.0 +croniter==1.2.0 # via # -c requirements.txt # flytekit @@ -79,11 +75,12 @@ cryptography==36.0.1 # via # -c requirements.txt # paramiko + # secretstorage dataclasses-json==0.5.6 # via # -c requirements.txt # flytekit -decorator==5.1.0 +decorator==5.1.1 # via # -c requirements.txt # retry @@ -91,7 +88,7 @@ deprecated==1.2.13 # via # -c requirements.txt # flytekit -diskcache==5.3.0 +diskcache==5.4.0 # via # -c requirements.txt # flytekit @@ -117,18 +114,18 @@ docstring-parser==0.13 # via # -c requirements.txt # flytekit -filelock==3.4.0 +filelock==3.4.2 # via virtualenv -flyteidl==0.21.17 +flyteidl==0.22.0 # via # -c requirements.txt # flytekit -google-api-core[grpc]==2.4.0 +google-api-core[grpc]==2.5.0 # via # google-cloud-bigquery # google-cloud-bigquery-storage # google-cloud-core -google-auth==2.3.3 +google-auth==2.6.0 # via # google-api-core # google-cloud-core @@ -140,7 +137,7 @@ google-cloud-core==2.2.2 # via google-cloud-bigquery google-crc32c==1.3.0 # via google-resumable-media -google-resumable-media==2.1.0 +google-resumable-media==2.2.0 # via google-cloud-bigquery googleapis-common-protos==1.54.0 # via @@ -155,18 +152,23 @@ grpcio==1.43.0 # grpcio-status grpcio-status==1.43.0 # via google-api-core -identify==2.4.0 +identify==2.4.8 # via pre-commit idna==3.3 # via # -c requirements.txt # requests -importlib-metadata==4.10.0 +importlib-metadata==4.10.1 # via # -c requirements.txt # keyring iniconfig==1.1.1 # via pytest +jeepney==0.7.1 + # via + # -c requirements.txt + # keyring + # secretstorage jinja2==3.0.3 # via # -c requirements.txt @@ -183,11 +185,11 @@ jsonschema==3.2.0 # via # -c requirements.txt # docker-compose -keyring==23.4.0 +keyring==23.5.0 # via # -c requirements.txt # flytekit -libcst==0.4.0 +libcst==0.4.1 # via google-cloud-bigquery-storage markupsafe==2.0.1 # via @@ -209,14 +211,14 @@ marshmallow-jsonschema==0.13.0 # flytekit mock==4.0.3 # via -r dev-requirements.in -mypy==0.930 +mypy==0.931 # via -r dev-requirements.in mypy-extensions==0.4.3 # via # -c requirements.txt # mypy # typing-inspect -natsort==8.0.2 +natsort==8.1.0 # via # -c requirements.txt # flytekit @@ -229,34 +231,29 @@ numpy==1.21.5 # pyarrow packaging==21.3 # via - # -c requirements.txt # google-cloud-bigquery # pytest pandas==1.3.5 # via # -c requirements.txt # flytekit -paramiko==2.8.1 - # via - # -c requirements.txt - # docker -platformdirs==2.4.0 - # via - # -c requirements.txt - # virtualenv +paramiko==2.9.2 + # via docker +platformdirs==2.4.1 + # via virtualenv pluggy==1.0.0 # via pytest poyo==0.5.0 # via # -c requirements.txt # cookiecutter -pre-commit==2.16.0 +pre-commit==2.17.0 # via -r dev-requirements.in -proto-plus==1.19.8 +proto-plus==1.20.0 # via # google-cloud-bigquery # google-cloud-bigquery-storage -protobuf==3.19.1 +protobuf==3.19.4 # via # -c requirements.txt # flyteidl @@ -285,15 +282,11 @@ pycparser==2.21 # via # -c requirements.txt # cffi -pynacl==1.4.0 - # via - # -c requirements.txt - # paramiko -pyparsing==3.0.6 - # via - # -c requirements.txt - # packaging -pyrsistent==0.18.0 +pynacl==1.5.0 + # via paramiko +pyparsing==3.0.7 + # via packaging +pyrsistent==0.18.1 # via # -c requirements.txt # jsonschema @@ -306,7 +299,7 @@ pytest-docker==0.10.3 # via pytest-flyte pytest-flyte @ git+https://github.com/flyteorg/pytest-flyte@main # via -r dev-requirements.in -python-dateutil==2.8.1 +python-dateutil==2.8.2 # via # -c requirements.txt # arrow @@ -339,11 +332,11 @@ pyyaml==5.4.1 # docker-compose # libcst # pre-commit -regex==2021.11.10 +regex==2022.1.18 # via # -c requirements.txt # docker-image-py -requests==2.26.0 +requests==2.27.1 # via # -c requirements.txt # cookiecutter @@ -353,7 +346,7 @@ requests==2.26.0 # google-api-core # google-cloud-bigquery # responses -responses==0.16.0 +responses==0.18.0 # via # -c requirements.txt # flytekit @@ -363,6 +356,10 @@ retry==0.9.2 # flytekit rsa==4.8 # via google-auth +secretstorage==3.3.1 + # via + # -c requirements.txt + # keyring six==1.16.0 # via # -c requirements.txt @@ -372,9 +369,7 @@ six==1.16.0 # google-auth # grpcio # jsonschema - # pynacl # python-dateutil - # responses # virtualenv # websocket-client sortedcontainers==2.4.0 @@ -395,14 +390,14 @@ toml==0.10.2 # via # pre-commit # pytest -tomli==1.2.3 +tomli==2.0.0 # via - # -c requirements.txt # coverage # mypy typing-extensions==4.0.1 # via # -c requirements.txt + # flytekit # libcst # mypy # typing-inspect @@ -411,13 +406,13 @@ typing-inspect==0.7.1 # -c requirements.txt # dataclasses-json # libcst -urllib3==1.26.7 +urllib3==1.26.8 # via # -c requirements.txt # flytekit # requests # responses -virtualenv==20.10.0 +virtualenv==20.13.1 # via pre-commit websocket-client==0.59.0 # via @@ -432,7 +427,7 @@ wrapt==1.13.3 # -c requirements.txt # deprecated # flytekit -zipp==3.6.0 +zipp==3.7.0 # via # -c requirements.txt # importlib-metadata diff --git a/doc-requirements.txt b/doc-requirements.txt index 9590c1face..c320d836d5 100644 --- a/doc-requirements.txt +++ b/doc-requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.8 +# This file is autogenerated by pip-compile with python 3.9 # To update, run: # # make doc-requirements.txt @@ -8,26 +8,12 @@ # via -r doc-requirements.in alabaster==0.7.12 # via sphinx -ansiwrap==0.8.4 - # via papermill -appnope==0.1.2 - # via - # ipykernel - # ipython -arrow==1.2.1 +arrow==1.2.2 # via jinja2-time astroid==2.9.3 # via sphinx-autoapi -asttokens==2.0.5 - # via stack-data -attrs==21.4.0 - # via jsonschema babel==2.9.1 # via sphinx -backcall==0.2.0 - # via ipython -bcrypt==3.2.0 - # via paramiko beautifulsoup4==4.10.0 # via # furo @@ -35,58 +21,36 @@ beautifulsoup4==4.10.0 # sphinx-material binaryornot==0.4.4 # via cookiecutter -black==21.12b0 - # via - # ipython - # papermill -bleach==4.1.0 - # via nbconvert -boto3==1.20.34 - # via sagemaker-training -botocore==1.23.34 - # via - # boto3 - # s3transfer certifi==2021.10.8 # via requests cffi==1.15.0 - # via - # bcrypt - # cryptography - # pynacl + # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.10 +charset-normalizer==2.0.11 # via requests checksumdir==1.2.0 # via flytekit -click==7.1.2 +click==8.0.3 # via - # black # cookiecutter # flytekit - # hmsclient - # papermill cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.1.0 +croniter==1.2.0 # via flytekit cryptography==36.0.1 # via # -r doc-requirements.in - # paramiko + # secretstorage css-html-js-minify==2.5.5 # via sphinx-material dataclasses-json==0.5.6 # via flytekit decorator==5.1.1 - # via - # ipython - # retry -defusedxml==0.7.1 - # via nbconvert + # via retry deprecated==1.2.13 # via flytekit diskcache==5.4.0 @@ -99,75 +63,34 @@ docutils==0.17.1 # via # sphinx # sphinx-panels -entrypoints==0.3 - # via - # jupyter-client - # nbconvert - # papermill -executing==0.8.2 - # via stack-data -flyteidl==0.21.22 +flyteidl==0.22.0 # via flytekit furo @ git+https://github.com/flyteorg/furo@main # via -r doc-requirements.in -gevent==21.12.0 - # via sagemaker-training -greenlet==1.1.2 - # via gevent grpcio==1.43.0 # via # -r doc-requirements.in # flytekit -hmsclient==0.1.1 - # via flytekit idna==3.3 # via requests imagesize==1.3.0 # via sphinx -importlib-metadata==4.10.0 - # via keyring -importlib-resources==5.4.0 - # via jsonschema -inotify_simple==1.2.1 - # via sagemaker-training -ipykernel==5.5.6 - # via flytekit -ipython==8.0.0 - # via ipykernel -ipython-genutils==0.2.0 +importlib-metadata==4.10.1 # via - # ipykernel - # nbformat -jedi==0.18.1 - # via ipython + # keyring + # sphinx +jeepney==0.7.1 + # via + # keyring + # secretstorage jinja2==3.0.3 # via # cookiecutter # jinja2-time - # nbconvert # sphinx # sphinx-autoapi jinja2-time==0.2.0 # via cookiecutter -jmespath==0.10.0 - # via - # boto3 - # botocore -jsonschema==4.4.0 - # via nbformat -jupyter-client==7.1.0 - # via - # ipykernel - # nbclient -jupyter-core==4.9.1 - # via - # jupyter-client - # nbconvert - # nbformat -jupyterlab-pygments==0.1.2 - # via nbconvert -k8s-proto==0.0.3 - # via flytekit keyring==23.5.0 # via flytekit lazy-object-proxy==1.7.1 @@ -185,106 +108,41 @@ marshmallow-enum==1.5.1 # via dataclasses-json marshmallow-jsonschema==0.13.0 # via flytekit -matplotlib-inline==0.1.3 - # via ipython -mistune==0.8.4 - # via nbconvert mypy-extensions==0.4.3 - # via - # black - # typing-inspect -natsort==8.0.2 - # via flytekit -nbclient==0.5.9 - # via - # nbconvert - # papermill -nbconvert==6.4.0 + # via typing-inspect +natsort==8.1.0 # via flytekit -nbformat==5.1.3 +numpy==1.22.2 # via - # nbclient - # nbconvert - # papermill -nest-asyncio==1.5.4 - # via - # jupyter-client - # nbclient -numpy==1.22.0 - # via - # flytekit # pandas # pyarrow - # sagemaker-training - # scipy packaging==21.3 - # via - # bleach - # sphinx -pandas==1.3.5 - # via flytekit -pandocfilters==1.5.0 - # via nbconvert -papermill==2.3.3 + # via sphinx +pandas==1.4.0 # via flytekit -paramiko==2.9.2 - # via sagemaker-training -parso==0.8.3 - # via jedi -pathspec==0.9.0 - # via black -pexpect==4.8.0 - # via ipython -pickleshare==0.7.5 - # via ipython -platformdirs==2.4.1 - # via black poyo==0.5.0 # via cookiecutter -prompt-toolkit==3.0.24 - # via ipython -protobuf==3.19.3 +protobuf==3.19.4 # via # flyteidl # flytekit - # k8s-proto - # sagemaker-training -psutil==5.9.0 - # via sagemaker-training -ptyprocess==0.7.0 - # via pexpect -pure-eval==0.2.1 - # via stack-data py==1.11.0 # via retry -py4j==0.10.9.2 - # via pyspark pyarrow==6.0.1 # via flytekit pycparser==2.21 # via cffi pygments==2.11.2 # via - # ipython - # jupyterlab-pygments - # nbconvert # sphinx # sphinx-prompt -pynacl==1.5.0 - # via paramiko -pyparsing==3.0.6 +pyparsing==3.0.7 # via packaging -pyrsistent==0.18.0 - # via jsonschema -pyspark==3.2.0 - # via flytekit -python-dateutil==2.8.1 +python-dateutil==2.8.2 # via # arrow - # botocore # croniter # flytekit - # jupyter-client # pandas python-json-logger==2.0.2 # via flytekit @@ -300,53 +158,34 @@ pytz==2021.3 # flytekit # pandas pyyaml==6.0 - # via - # papermill - # sphinx-autoapi -pyzmq==22.3.0 - # via jupyter-client -regex==2021.11.10 + # via sphinx-autoapi +regex==2022.1.18 # via docker-image-py requests==2.27.1 # via # cookiecutter # flytekit - # papermill # responses # sphinx -responses==0.17.0 +responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit -retrying==1.3.3 - # via sagemaker-training -s3transfer==0.5.0 - # via boto3 -sagemaker-training==3.9.2 - # via flytekit -scipy==1.7.3 - # via sagemaker-training +secretstorage==3.3.1 + # via keyring six==1.16.0 # via - # asttokens - # bcrypt - # bleach # cookiecutter - # flytekit # grpcio # python-dateutil - # responses - # retrying - # sagemaker-training # sphinx-code-include - # thrift snowballstemmer==2.2.0 # via sphinx sortedcontainers==2.4.0 # via flytekit soupsieve==2.3.1 # via beautifulsoup4 -sphinx==4.3.2 +sphinx==4.4.0 # via # -r doc-requirements.in # furo @@ -363,7 +202,7 @@ sphinx-autoapi==1.8.4 # via -r doc-requirements.in sphinx-code-include==1.1.1 # via -r doc-requirements.in -sphinx-copybutton==0.4.0 +sphinx-copybutton==0.5.0 # via -r doc-requirements.in sphinx-fontawesome==0.0.6 # via -r doc-requirements.in @@ -389,42 +228,14 @@ sphinxcontrib-serializinghtml==1.1.5 # via sphinx sphinxcontrib-yt==0.2.2 # via -r doc-requirements.in -stack-data==0.1.3 - # via ipython statsd==3.3.0 # via flytekit -tenacity==8.0.1 - # via papermill -testpath==0.5.0 - # via nbconvert text-unidecode==1.3 # via python-slugify -textwrap3==0.9.2 - # via ansiwrap -thrift==0.15.0 - # via hmsclient -tomli==1.2.3 - # via black -tornado==6.1 - # via - # ipykernel - # jupyter-client -tqdm==4.62.3 - # via papermill -traitlets==5.1.1 - # via - # ipykernel - # ipython - # jupyter-client - # jupyter-core - # matplotlib-inline - # nbclient - # nbconvert - # nbformat typing-extensions==4.0.1 # via # astroid - # black + # flytekit # typing-inspect typing-inspect==0.7.1 # via dataclasses-json @@ -434,16 +245,9 @@ unidecode==1.3.2 # sphinx-autoapi urllib3==1.26.8 # via - # botocore # flytekit # requests # responses -wcwidth==0.2.5 - # via prompt-toolkit -webencodings==0.5.1 - # via bleach -werkzeug==2.0.2 - # via sagemaker-training wheel==0.37.1 # via flytekit wrapt==1.13.3 @@ -452,14 +256,7 @@ wrapt==1.13.3 # deprecated # flytekit zipp==3.7.0 - # via - # importlib-metadata - # importlib-resources -zope.event==4.5.0 - # via gevent -zope.interface==5.4.0 - # via gevent + # via importlib-metadata # The following packages are considered to be unsafe in a requirements file: -# pip # setuptools diff --git a/plugins/flytekit-aws-athena/requirements.txt b/plugins/flytekit-aws-athena/requirements.txt index 4cd07971c4..6fe91c53c8 100644 --- a/plugins/flytekit-aws-athena/requirements.txt +++ b/plugins/flytekit-aws-athena/requirements.txt @@ -6,7 +6,7 @@ # -e file:.#egg=flytekitplugins-athena # via -r requirements.in -arrow==1.2.1 +arrow==1.2.2 # via jinja2-time binaryornot==0.4.4 # via cookiecutter @@ -16,7 +16,7 @@ cffi==1.15.0 # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.7 +charset-normalizer==2.0.11 # via requests checksumdir==1.2.0 # via flytekit @@ -28,31 +28,31 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.0.15 +croniter==1.2.0 # via flytekit -cryptography==35.0.0 +cryptography==36.0.1 # via secretstorage dataclasses-json==0.5.6 # via flytekit -decorator==5.1.0 +decorator==5.1.1 # via retry deprecated==1.2.13 # via flytekit -diskcache==5.2.1 +diskcache==5.4.0 # via flytekit docker-image-py==0.1.12 # via flytekit -docstring-parser==0.12 +docstring-parser==0.13 # via flytekit -flyteidl==0.21.8 +flyteidl==0.22.0 # via flytekit -flytekit==0.24.0 +flytekit==0.30.0 # via flytekitplugins-athena -grpcio==1.41.1 +grpcio==1.43.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.8.2 +importlib-metadata==4.10.1 # via keyring jeepney==0.7.1 # via @@ -64,11 +64,11 @@ jinja2==3.0.3 # jinja2-time jinja2-time==0.2.0 # via cookiecutter -keyring==23.2.1 +keyring==23.5.0 # via flytekit markupsafe==2.0.1 # via jinja2 -marshmallow==3.14.0 +marshmallow==3.14.1 # via # dataclasses-json # marshmallow-enum @@ -79,27 +79,27 @@ marshmallow-jsonschema==0.13.0 # via flytekit mypy-extensions==0.4.3 # via typing-inspect -natsort==8.0.0 +natsort==8.1.0 # via flytekit -numpy==1.21.4 +numpy==1.22.2 # via # pandas # pyarrow -pandas==1.3.4 +pandas==1.4.0 # via flytekit poyo==0.5.0 # via cookiecutter -protobuf==3.19.1 +protobuf==3.19.4 # via # flyteidl # flytekit py==1.11.0 # via retry -pyarrow==6.0.0 +pyarrow==6.0.1 # via flytekit pycparser==2.21 # via cffi -python-dateutil==2.8.1 +python-dateutil==2.8.2 # via # arrow # croniter @@ -111,18 +111,18 @@ python-slugify==5.0.2 # via cookiecutter pytimeparse==1.1.8 # via flytekit -pytz==2018.4 +pytz==2021.3 # via # flytekit # pandas -regex==2021.11.10 +regex==2022.1.18 # via docker-image-py -requests==2.26.0 +requests==2.27.1 # via # cookiecutter # flytekit # responses -responses==0.15.0 +responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit @@ -131,30 +131,30 @@ secretstorage==3.3.1 six==1.16.0 # via # cookiecutter - # flytekit # grpcio # python-dateutil - # responses sortedcontainers==2.4.0 # via flytekit statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -typing-extensions==3.10.0.2 - # via typing-inspect +typing-extensions==4.0.1 + # via + # flytekit + # typing-inspect typing-inspect==0.7.1 # via dataclasses-json -urllib3==1.26.7 +urllib3==1.26.8 # via # flytekit # requests # responses -wheel==0.37.0 +wheel==0.37.1 # via flytekit wrapt==1.13.3 # via # deprecated # flytekit -zipp==3.6.0 +zipp==3.7.0 # via importlib-metadata diff --git a/plugins/flytekit-aws-sagemaker/requirements.txt b/plugins/flytekit-aws-sagemaker/requirements.txt index b5d42c422d..2624f28a98 100644 --- a/plugins/flytekit-aws-sagemaker/requirements.txt +++ b/plugins/flytekit-aws-sagemaker/requirements.txt @@ -6,15 +6,15 @@ # -e file:.#egg=flytekitplugins-awssagemaker # via -r requirements.in -arrow==1.2.1 +arrow==1.2.2 # via jinja2-time bcrypt==3.2.0 # via paramiko binaryornot==0.4.4 # via cookiecutter -boto3==1.20.3 +boto3==1.20.50 # via sagemaker-training -botocore==1.23.3 +botocore==1.23.50 # via # boto3 # s3transfer @@ -27,7 +27,7 @@ cffi==1.15.0 # pynacl chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.7 +charset-normalizer==2.0.11 # via requests checksumdir==1.2.0 # via flytekit @@ -39,39 +39,39 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.0.15 +croniter==1.2.0 # via flytekit -cryptography==35.0.0 +cryptography==36.0.1 # via # paramiko # secretstorage dataclasses-json==0.5.6 # via flytekit -decorator==5.1.0 +decorator==5.1.1 # via retry deprecated==1.2.13 # via flytekit -diskcache==5.2.1 +diskcache==5.4.0 # via flytekit docker-image-py==0.1.12 # via flytekit -docstring-parser==0.12 +docstring-parser==0.13 # via flytekit -flyteidl==0.21.8 +flyteidl==0.22.0 # via flytekit -flytekit==0.24.0 +flytekit==0.30.0 # via flytekitplugins-awssagemaker -gevent==21.8.0 +gevent==21.12.0 # via sagemaker-training greenlet==1.1.2 # via gevent -grpcio==1.41.1 +grpcio==1.43.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.8.2 +importlib-metadata==4.10.1 # via keyring -inotify_simple==1.2.1 +inotify-simple==1.2.1 # via sagemaker-training jeepney==0.7.1 # via @@ -87,11 +87,11 @@ jmespath==0.10.0 # via # boto3 # botocore -keyring==23.2.1 +keyring==23.5.0 # via flytekit markupsafe==2.0.1 # via jinja2 -marshmallow==3.14.0 +marshmallow==3.14.1 # via # dataclasses-json # marshmallow-enum @@ -102,36 +102,36 @@ marshmallow-jsonschema==0.13.0 # via flytekit mypy-extensions==0.4.3 # via typing-inspect -natsort==8.0.0 +natsort==8.1.0 # via flytekit -numpy==1.21.4 +numpy==1.22.2 # via # pandas # pyarrow # sagemaker-training # scipy -pandas==1.3.4 +pandas==1.4.0 # via flytekit -paramiko==2.8.0 +paramiko==2.9.2 # via sagemaker-training poyo==0.5.0 # via cookiecutter -protobuf==3.19.1 +protobuf==3.19.4 # via # flyteidl # flytekit # sagemaker-training -psutil==5.8.0 +psutil==5.9.0 # via sagemaker-training py==1.11.0 # via retry -pyarrow==6.0.0 +pyarrow==6.0.1 # via flytekit pycparser==2.21 # via cffi -pynacl==1.4.0 +pynacl==1.5.0 # via paramiko -python-dateutil==2.8.1 +python-dateutil==2.8.2 # via # arrow # botocore @@ -144,28 +144,28 @@ python-slugify==5.0.2 # via cookiecutter pytimeparse==1.1.8 # via flytekit -pytz==2018.4 +pytz==2021.3 # via # flytekit # pandas -regex==2021.11.10 +regex==2022.1.18 # via docker-image-py -requests==2.26.0 +requests==2.27.1 # via # cookiecutter # flytekit # responses -responses==0.15.0 +responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit retrying==1.3.3 # via sagemaker-training -s3transfer==0.5.0 +s3transfer==0.5.1 # via boto3 sagemaker-training==3.9.2 # via flytekitplugins-awssagemaker -scipy==1.7.2 +scipy==1.8.0 # via sagemaker-training secretstorage==3.3.1 # via keyring @@ -173,11 +173,8 @@ six==1.16.0 # via # bcrypt # cookiecutter - # flytekit # grpcio - # pynacl # python-dateutil - # responses # retrying # sagemaker-training sortedcontainers==2.4.0 @@ -186,29 +183,31 @@ statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -typing-extensions==3.10.0.2 - # via typing-inspect +typing-extensions==4.0.1 + # via + # flytekit + # typing-inspect typing-inspect==0.7.1 # via dataclasses-json -urllib3==1.26.7 +urllib3==1.26.8 # via # botocore # flytekit # requests # responses -werkzeug==2.0.2 +werkzeug==2.0.3 # via sagemaker-training -wheel==0.37.0 +wheel==0.37.1 # via flytekit wrapt==1.13.3 # via # deprecated # flytekit -zipp==3.6.0 +zipp==3.7.0 # via importlib-metadata -zope.event==4.5.0 +zope-event==4.5.0 # via gevent -zope.interface==5.4.0 +zope-interface==5.4.0 # via gevent # The following packages are considered to be unsafe in a requirements file: diff --git a/plugins/flytekit-bigquery/requirements.txt b/plugins/flytekit-bigquery/requirements.txt index 45be4f75a0..f0bc647453 100644 --- a/plugins/flytekit-bigquery/requirements.txt +++ b/plugins/flytekit-bigquery/requirements.txt @@ -6,17 +6,19 @@ # -e file:.#egg=flytekitplugins-bigquery # via -r requirements.in -arrow==1.2.1 +arrow==1.2.2 # via jinja2-time binaryornot==0.4.4 # via cookiecutter -cachetools==4.2.4 +cachetools==5.0.0 # via google-auth certifi==2021.10.8 # via requests +cffi==1.15.0 + # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.7 +charset-normalizer==2.0.11 # via requests checksumdir==1.2.0 # via flytekit @@ -28,39 +30,41 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.0.15 +croniter==1.2.0 # via flytekit +cryptography==36.0.1 + # via secretstorage dataclasses-json==0.5.6 # via flytekit -decorator==5.1.0 +decorator==5.1.1 # via retry deprecated==1.2.13 # via flytekit -diskcache==5.2.1 +diskcache==5.4.0 # via flytekit docker-image-py==0.1.12 # via flytekit -docstring-parser==0.12 +docstring-parser==0.13 # via flytekit -flyteidl==0.21.8 +flyteidl==0.22.0 # via flytekit -flytekit==0.24.0 +flytekit==0.30.0 # via flytekitplugins-bigquery -google-api-core[grpc]==2.3.2 +google-api-core[grpc]==2.5.0 # via # google-cloud-bigquery # google-cloud-core -google-auth==2.3.3 +google-auth==2.6.0 # via # google-api-core # google-cloud-core -google-cloud-bigquery==2.31.0 +google-cloud-bigquery==2.32.0 # via flytekitplugins-bigquery -google-cloud-core==2.2.1 +google-cloud-core==2.2.2 # via google-cloud-bigquery google-crc32c==1.3.0 # via google-resumable-media -google-resumable-media==2.1.0 +google-resumable-media==2.2.0 # via google-cloud-bigquery googleapis-common-protos==1.54.0 # via @@ -76,19 +80,23 @@ grpcio-status==1.43.0 # via google-api-core idna==3.3 # via requests -importlib-metadata==4.8.2 +importlib-metadata==4.10.1 # via keyring +jeepney==0.7.1 + # via + # keyring + # secretstorage jinja2==3.0.3 # via # cookiecutter # jinja2-time jinja2-time==0.2.0 # via cookiecutter -keyring==23.2.1 +keyring==23.5.0 # via flytekit markupsafe==2.0.1 # via jinja2 -marshmallow==3.14.0 +marshmallow==3.14.1 # via # dataclasses-json # marshmallow-enum @@ -99,21 +107,21 @@ marshmallow-jsonschema==0.13.0 # via flytekit mypy-extensions==0.4.3 # via typing-inspect -natsort==8.0.0 +natsort==8.1.0 # via flytekit -numpy==1.21.4 +numpy==1.22.2 # via # pandas # pyarrow packaging==21.3 # via google-cloud-bigquery -pandas==1.3.4 +pandas==1.4.0 # via flytekit poyo==0.5.0 # via cookiecutter -proto-plus==1.19.8 +proto-plus==1.20.0 # via google-cloud-bigquery -protobuf==3.19.1 +protobuf==3.19.4 # via # flyteidl # flytekit @@ -124,7 +132,7 @@ protobuf==3.19.1 # proto-plus py==1.11.0 # via retry -pyarrow==6.0.0 +pyarrow==6.0.1 # via flytekit pyasn1==0.4.8 # via @@ -132,9 +140,11 @@ pyasn1==0.4.8 # rsa pyasn1-modules==0.2.8 # via google-auth -pyparsing==3.0.6 +pycparser==2.21 + # via cffi +pyparsing==3.0.7 # via packaging -python-dateutil==2.8.1 +python-dateutil==2.8.2 # via # arrow # croniter @@ -147,56 +157,55 @@ python-slugify==5.0.2 # via cookiecutter pytimeparse==1.1.8 # via flytekit -pytz==2018.4 +pytz==2021.3 # via # flytekit # pandas -regex==2021.11.10 +regex==2022.1.18 # via docker-image-py -requests==2.26.0 +requests==2.27.1 # via # cookiecutter # flytekit # google-api-core # google-cloud-bigquery # responses -responses==0.15.0 +responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit rsa==4.8 # via google-auth +secretstorage==3.3.1 + # via keyring six==1.16.0 # via # cookiecutter - # flytekit # google-auth # grpcio # python-dateutil - # responses sortedcontainers==2.4.0 # via flytekit statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -typing-extensions==3.10.0.2 - # via typing-inspect +typing-extensions==4.0.1 + # via + # flytekit + # typing-inspect typing-inspect==0.7.1 # via dataclasses-json -urllib3==1.26.7 +urllib3==1.26.8 # via # flytekit # requests # responses -wheel==0.37.0 +wheel==0.37.1 # via flytekit wrapt==1.13.3 # via # deprecated # flytekit -zipp==3.6.0 +zipp==3.7.0 # via importlib-metadata - -# The following packages are considered to be unsafe in a requirements file: -# setuptools diff --git a/plugins/flytekit-data-fsspec/requirements.txt b/plugins/flytekit-data-fsspec/requirements.txt index d6ee7128aa..a14a5bf6cf 100644 --- a/plugins/flytekit-data-fsspec/requirements.txt +++ b/plugins/flytekit-data-fsspec/requirements.txt @@ -6,7 +6,7 @@ # -e file:.#egg=flytekitplugins-data-fsspec # via -r requirements.in -arrow==1.2.1 +arrow==1.2.2 # via jinja2-time binaryornot==0.4.4 # via cookiecutter @@ -16,7 +16,7 @@ cffi==1.15.0 # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.7 +charset-normalizer==2.0.11 # via requests checksumdir==1.2.0 # via flytekit @@ -28,33 +28,33 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.0.15 +croniter==1.2.0 # via flytekit -cryptography==35.0.0 +cryptography==36.0.1 # via secretstorage dataclasses-json==0.5.6 # via flytekit -decorator==5.1.0 +decorator==5.1.1 # via retry deprecated==1.2.13 # via flytekit -diskcache==5.2.1 +diskcache==5.4.0 # via flytekit docker-image-py==0.1.12 # via flytekit -docstring-parser==0.12 +docstring-parser==0.13 # via flytekit -flyteidl==0.21.8 +flyteidl==0.22.0 # via flytekit -flytekit==0.24.0 +flytekit==0.30.0 # via flytekitplugins-data-fsspec -fsspec==2021.11.0 +fsspec==2022.1.0 # via flytekitplugins-data-fsspec -grpcio==1.41.1 +grpcio==1.43.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.8.2 +importlib-metadata==4.10.1 # via keyring jeepney==0.7.1 # via @@ -66,11 +66,11 @@ jinja2==3.0.3 # jinja2-time jinja2-time==0.2.0 # via cookiecutter -keyring==23.2.1 +keyring==23.5.0 # via flytekit markupsafe==2.0.1 # via jinja2 -marshmallow==3.14.0 +marshmallow==3.14.1 # via # dataclasses-json # marshmallow-enum @@ -81,27 +81,27 @@ marshmallow-jsonschema==0.13.0 # via flytekit mypy-extensions==0.4.3 # via typing-inspect -natsort==8.0.0 +natsort==8.1.0 # via flytekit -numpy==1.21.4 +numpy==1.22.2 # via # pandas # pyarrow -pandas==1.3.4 +pandas==1.4.0 # via flytekit poyo==0.5.0 # via cookiecutter -protobuf==3.19.1 +protobuf==3.19.4 # via # flyteidl # flytekit py==1.11.0 # via retry -pyarrow==6.0.0 +pyarrow==6.0.1 # via flytekit pycparser==2.21 # via cffi -python-dateutil==2.8.1 +python-dateutil==2.8.2 # via # arrow # croniter @@ -113,18 +113,18 @@ python-slugify==5.0.2 # via cookiecutter pytimeparse==1.1.8 # via flytekit -pytz==2018.4 +pytz==2021.3 # via # flytekit # pandas -regex==2021.11.10 +regex==2022.1.18 # via docker-image-py -requests==2.26.0 +requests==2.27.1 # via # cookiecutter # flytekit # responses -responses==0.15.0 +responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit @@ -133,30 +133,30 @@ secretstorage==3.3.1 six==1.16.0 # via # cookiecutter - # flytekit # grpcio # python-dateutil - # responses sortedcontainers==2.4.0 # via flytekit statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -typing-extensions==3.10.0.2 - # via typing-inspect +typing-extensions==4.0.1 + # via + # flytekit + # typing-inspect typing-inspect==0.7.1 # via dataclasses-json -urllib3==1.26.7 +urllib3==1.26.8 # via # flytekit # requests # responses -wheel==0.37.0 +wheel==0.37.1 # via flytekit wrapt==1.13.3 # via # deprecated # flytekit -zipp==3.6.0 +zipp==3.7.0 # via importlib-metadata diff --git a/plugins/flytekit-dolt/requirements.txt b/plugins/flytekit-dolt/requirements.txt index f669ec8672..559d5b2085 100644 --- a/plugins/flytekit-dolt/requirements.txt +++ b/plugins/flytekit-dolt/requirements.txt @@ -6,7 +6,7 @@ # -e file:.#egg=flytekitplugins-dolt # via -r requirements.in -arrow==1.2.1 +arrow==1.2.2 # via jinja2-time binaryornot==0.4.4 # via cookiecutter @@ -16,7 +16,7 @@ cffi==1.15.0 # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.7 +charset-normalizer==2.0.11 # via requests checksumdir==1.2.0 # via flytekit @@ -28,37 +28,37 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.0.15 +croniter==1.2.0 # via flytekit -cryptography==35.0.0 +cryptography==36.0.1 # via secretstorage dataclasses-json==0.5.6 # via # dolt-integrations # flytekit -decorator==5.1.0 +decorator==5.1.1 # via retry deprecated==1.2.13 # via flytekit -diskcache==5.2.1 +diskcache==5.4.0 # via flytekit docker-image-py==0.1.12 # via flytekit -docstring-parser==0.12 +docstring-parser==0.13 # via flytekit dolt-integrations==0.1.5 # via flytekitplugins-dolt -doltcli==0.1.15 +doltcli==0.1.17 # via dolt-integrations -flyteidl==0.21.8 +flyteidl==0.22.0 # via flytekit -flytekit==0.24.0 +flytekit==0.30.0 # via flytekitplugins-dolt -grpcio==1.41.1 +grpcio==1.43.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.8.2 +importlib-metadata==4.10.1 # via keyring jeepney==0.7.1 # via @@ -70,11 +70,11 @@ jinja2==3.0.3 # jinja2-time jinja2-time==0.2.0 # via cookiecutter -keyring==23.2.1 +keyring==23.5.0 # via flytekit markupsafe==2.0.1 # via jinja2 -marshmallow==3.14.0 +marshmallow==3.14.1 # via # dataclasses-json # marshmallow-enum @@ -85,29 +85,29 @@ marshmallow-jsonschema==0.13.0 # via flytekit mypy-extensions==0.4.3 # via typing-inspect -natsort==8.0.0 +natsort==8.1.0 # via flytekit -numpy==1.21.4 +numpy==1.22.2 # via # pandas # pyarrow -pandas==1.3.4 +pandas==1.4.0 # via # dolt-integrations # flytekit poyo==0.5.0 # via cookiecutter -protobuf==3.19.1 +protobuf==3.19.4 # via # flyteidl # flytekit py==1.11.0 # via retry -pyarrow==6.0.0 +pyarrow==6.0.1 # via flytekit pycparser==2.21 # via cffi -python-dateutil==2.8.1 +python-dateutil==2.8.2 # via # arrow # croniter @@ -119,18 +119,18 @@ python-slugify==5.0.2 # via cookiecutter pytimeparse==1.1.8 # via flytekit -pytz==2018.4 +pytz==2021.3 # via # flytekit # pandas -regex==2021.11.10 +regex==2022.1.18 # via docker-image-py -requests==2.26.0 +requests==2.27.1 # via # cookiecutter # flytekit # responses -responses==0.15.0 +responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit @@ -139,30 +139,30 @@ secretstorage==3.3.1 six==1.16.0 # via # cookiecutter - # flytekit # grpcio # python-dateutil - # responses sortedcontainers==2.4.0 # via flytekit statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -typing-extensions==3.10.0.2 - # via typing-inspect +typing-extensions==4.0.1 + # via + # flytekit + # typing-inspect typing-inspect==0.7.1 # via dataclasses-json -urllib3==1.26.7 +urllib3==1.26.8 # via # flytekit # requests # responses -wheel==0.37.0 +wheel==0.37.1 # via flytekit wrapt==1.13.3 # via # deprecated # flytekit -zipp==3.6.0 +zipp==3.7.0 # via importlib-metadata diff --git a/plugins/flytekit-greatexpectations/requirements.txt b/plugins/flytekit-greatexpectations/requirements.txt index 5bdbc66553..fdbb857264 100644 --- a/plugins/flytekit-greatexpectations/requirements.txt +++ b/plugins/flytekit-greatexpectations/requirements.txt @@ -6,34 +6,41 @@ # -e file:.#egg=flytekitplugins-great_expectations # via -r requirements.in -altair==4.1.0 +altair==4.2.0 # via great-expectations -argon2-cffi==21.1.0 +argon2-cffi==21.3.0 # via notebook -arrow==1.2.1 +argon2-cffi-bindings==21.2.0 + # via argon2-cffi +arrow==1.2.2 # via jinja2-time -attrs==21.2.0 +asttokens==2.0.5 + # via stack-data +attrs==21.4.0 # via jsonschema backcall==0.2.0 # via ipython binaryornot==0.4.4 # via cookiecutter +black==21.12b0 + # via ipython bleach==4.1.0 # via nbconvert certifi==2021.10.8 # via requests cffi==1.15.0 # via - # argon2-cffi + # argon2-cffi-bindings # cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.7 +charset-normalizer==2.0.10 # via requests checksumdir==1.2.0 # via flytekit click==7.1.2 # via + # black # cookiecutter # flytekit # great-expectations @@ -41,15 +48,15 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.0.15 +croniter==1.2.0 # via flytekit -cryptography==35.0.0 +cryptography==36.0.1 # via secretstorage dataclasses-json==0.5.6 # via flytekit debugpy==1.5.1 # via ipykernel -decorator==5.1.0 +decorator==5.1.1 # via # ipython # retry @@ -57,38 +64,40 @@ defusedxml==0.7.1 # via nbconvert deprecated==1.2.13 # via flytekit -diskcache==5.2.1 +diskcache==5.4.0 # via flytekit docker-image-py==0.1.12 # via flytekit -docstring-parser==0.12 +docstring-parser==0.13 # via flytekit entrypoints==0.3 # via # altair # jupyter-client # nbconvert -flyteidl==0.21.8 +executing==0.8.2 + # via stack-data +flyteidl==0.21.24 # via flytekit -flytekit==0.24.0 +flytekit==0.26.1 # via flytekitplugins-great-expectations -great-expectations==0.13.41 +great-expectations==0.14.3 # via flytekitplugins-great-expectations greenlet==1.1.2 # via sqlalchemy -grpcio==1.41.1 +grpcio==1.43.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.8.2 +importlib-metadata==4.10.1 # via # great-expectations # keyring -ipykernel==6.5.0 +ipykernel==6.7.0 # via # ipywidgets # notebook -ipython==7.31.1 +ipython==8.0.1 # via # ipykernel # ipywidgets @@ -99,7 +108,7 @@ ipython-genutils==0.2.0 # notebook ipywidgets==7.6.5 # via great-expectations -jedi==0.18.0 +jedi==0.18.1 # via ipython jeepney==0.7.1 # via @@ -119,12 +128,12 @@ jsonpatch==1.32 # via great-expectations jsonpointer==2.2 # via jsonpatch -jsonschema==4.2.1 +jsonschema==4.4.0 # via # altair # great-expectations # nbformat -jupyter-client==7.0.6 +jupyter-client==7.1.2 # via # ipykernel # nbclient @@ -139,11 +148,11 @@ jupyterlab-pygments==0.1.2 # via nbconvert jupyterlab-widgets==1.0.2 # via ipywidgets -keyring==23.2.1 +keyring==23.5.0 # via flytekit markupsafe==2.0.1 # via jinja2 -marshmallow==3.14.0 +marshmallow==3.14.1 # via # dataclasses-json # marshmallow-enum @@ -161,12 +170,14 @@ mistune==0.8.4 # great-expectations # nbconvert mypy-extensions==0.4.3 - # via typing-inspect -natsort==8.0.0 + # via + # black + # typing-inspect +natsort==8.0.2 # via flytekit -nbclient==0.5.5 +nbclient==0.5.10 # via nbconvert -nbconvert==6.2.0 +nbconvert==6.4.1 # via notebook nbformat==5.1.3 # via @@ -174,41 +185,47 @@ nbformat==5.1.3 # nbclient # nbconvert # notebook -nest-asyncio==1.5.1 +nest-asyncio==1.5.4 # via + # ipykernel # jupyter-client # nbclient -notebook==6.4.5 + # notebook +notebook==6.4.8 # via widgetsnbextension -numpy==1.21.4 +numpy==1.22.1 # via # altair # great-expectations # pandas # pyarrow # scipy -packaging==21.2 +packaging==21.3 # via bleach -pandas==1.3.4 +pandas==1.4.0 # via # altair # flytekit # great-expectations pandocfilters==1.5.0 # via nbconvert -parso==0.8.2 +parso==0.8.3 # via jedi +pathspec==0.9.0 + # via black pexpect==4.8.0 # via ipython pickleshare==0.7.5 # via ipython +platformdirs==2.4.1 + # via black poyo==0.5.0 # via cookiecutter -prometheus-client==0.12.0 +prometheus-client==0.13.1 # via notebook -prompt-toolkit==3.0.22 +prompt-toolkit==3.0.26 # via ipython -protobuf==3.19.1 +protobuf==3.19.4 # via # flyteidl # flytekit @@ -216,13 +233,15 @@ ptyprocess==0.7.0 # via # pexpect # terminado +pure-eval==0.2.2 + # via stack-data py==1.11.0 # via retry -pyarrow==6.0.0 +pyarrow==6.0.1 # via flytekit pycparser==2.21 # via cffi -pygments==2.10.0 +pygments==2.11.2 # via # ipython # jupyterlab-pygments @@ -231,7 +250,7 @@ pyparsing==2.4.7 # via # great-expectations # packaging -pyrsistent==0.18.0 +pyrsistent==0.18.1 # via jsonschema python-dateutil==2.8.1 # via @@ -247,7 +266,7 @@ python-slugify==5.0.2 # via cookiecutter pytimeparse==1.1.8 # via flytekit -pytz==2018.4 +pytz==2021.3 # via # flytekit # great-expectations @@ -258,15 +277,15 @@ pyzmq==22.3.0 # via # jupyter-client # notebook -regex==2021.11.10 +regex==2022.1.18 # via docker-image-py -requests==2.26.0 +requests==2.27.1 # via # cookiecutter # flytekit # great-expectations # responses -responses==0.15.0 +responses==0.17.0 # via flytekit retry==0.9.2 # via flytekit @@ -274,7 +293,7 @@ ruamel.yaml==0.17.17 # via great-expectations ruamel.yaml.clib==0.2.6 # via ruamel.yaml -scipy==1.7.2 +scipy==1.7.3 # via great-expectations secretstorage==3.3.1 # via keyring @@ -282,6 +301,7 @@ send2trash==1.8.0 # via notebook six==1.16.0 # via + # asttokens # bleach # cookiecutter # flytekit @@ -290,20 +310,24 @@ six==1.16.0 # responses sortedcontainers==2.4.0 # via flytekit -sqlalchemy==1.4.26 +sqlalchemy==1.4.31 # via # -r requirements.in # flytekitplugins-great-expectations +stack-data==0.1.4 + # via ipython statsd==3.3.0 # via flytekit termcolor==1.1.0 # via great-expectations -terminado==0.12.1 +terminado==0.13.1 # via notebook testpath==0.5.0 # via nbconvert text-unidecode==1.3 # via python-slugify +tomli==1.2.3 + # via black toolz==0.11.2 # via altair tornado==6.1 @@ -326,15 +350,17 @@ traitlets==5.1.1 # nbconvert # nbformat # notebook -typing-extensions==3.10.0.2 - # via typing-inspect +typing-extensions==4.0.1 + # via + # black + # typing-inspect typing-inspect==0.7.1 # via dataclasses-json tzdata==2021.5 # via pytz-deprecation-shim tzlocal==4.1 # via great-expectations -urllib3==1.26.7 +urllib3==1.26.8 # via # flytekit # requests @@ -343,7 +369,7 @@ wcwidth==0.2.5 # via prompt-toolkit webencodings==0.5.1 # via bleach -wheel==0.37.0 +wheel==0.37.1 # via flytekit widgetsnbextension==3.5.2 # via ipywidgets @@ -351,7 +377,7 @@ wrapt==1.13.3 # via # deprecated # flytekit -zipp==3.6.0 +zipp==3.7.0 # via importlib-metadata # The following packages are considered to be unsafe in a requirements file: diff --git a/plugins/flytekit-hive/requirements.txt b/plugins/flytekit-hive/requirements.txt index fe120d6f80..269d52fceb 100644 --- a/plugins/flytekit-hive/requirements.txt +++ b/plugins/flytekit-hive/requirements.txt @@ -6,7 +6,7 @@ # -e file:.#egg=flytekitplugins-hive # via -r requirements.in -arrow==1.2.1 +arrow==1.2.2 # via jinja2-time binaryornot==0.4.4 # via cookiecutter @@ -16,7 +16,7 @@ cffi==1.15.0 # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.7 +charset-normalizer==2.0.11 # via requests checksumdir==1.2.0 # via flytekit @@ -28,31 +28,31 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.0.15 +croniter==1.2.0 # via flytekit -cryptography==35.0.0 +cryptography==36.0.1 # via secretstorage dataclasses-json==0.5.6 # via flytekit -decorator==5.1.0 +decorator==5.1.1 # via retry deprecated==1.2.13 # via flytekit -diskcache==5.2.1 +diskcache==5.4.0 # via flytekit docker-image-py==0.1.12 # via flytekit -docstring-parser==0.12 +docstring-parser==0.13 # via flytekit -flyteidl==0.21.8 +flyteidl==0.22.0 # via flytekit -flytekit==0.24.0 +flytekit==0.30.0 # via flytekitplugins-hive -grpcio==1.41.1 +grpcio==1.43.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.8.2 +importlib-metadata==4.10.1 # via keyring jeepney==0.7.1 # via @@ -64,11 +64,11 @@ jinja2==3.0.3 # jinja2-time jinja2-time==0.2.0 # via cookiecutter -keyring==23.2.1 +keyring==23.5.0 # via flytekit markupsafe==2.0.1 # via jinja2 -marshmallow==3.14.0 +marshmallow==3.14.1 # via # dataclasses-json # marshmallow-enum @@ -79,27 +79,27 @@ marshmallow-jsonschema==0.13.0 # via flytekit mypy-extensions==0.4.3 # via typing-inspect -natsort==8.0.0 +natsort==8.1.0 # via flytekit -numpy==1.21.4 +numpy==1.22.2 # via # pandas # pyarrow -pandas==1.3.4 +pandas==1.4.0 # via flytekit poyo==0.5.0 # via cookiecutter -protobuf==3.19.1 +protobuf==3.19.4 # via # flyteidl # flytekit py==1.11.0 # via retry -pyarrow==6.0.0 +pyarrow==6.0.1 # via flytekit pycparser==2.21 # via cffi -python-dateutil==2.8.1 +python-dateutil==2.8.2 # via # arrow # croniter @@ -111,18 +111,18 @@ python-slugify==5.0.2 # via cookiecutter pytimeparse==1.1.8 # via flytekit -pytz==2018.4 +pytz==2021.3 # via # flytekit # pandas -regex==2021.11.10 +regex==2022.1.18 # via docker-image-py -requests==2.26.0 +requests==2.27.1 # via # cookiecutter # flytekit # responses -responses==0.15.0 +responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit @@ -131,30 +131,30 @@ secretstorage==3.3.1 six==1.16.0 # via # cookiecutter - # flytekit # grpcio # python-dateutil - # responses sortedcontainers==2.4.0 # via flytekit statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -typing-extensions==3.10.0.2 - # via typing-inspect +typing-extensions==4.0.1 + # via + # flytekit + # typing-inspect typing-inspect==0.7.1 # via dataclasses-json -urllib3==1.26.7 +urllib3==1.26.8 # via # flytekit # requests # responses -wheel==0.37.0 +wheel==0.37.1 # via flytekit wrapt==1.13.3 # via # deprecated # flytekit -zipp==3.6.0 +zipp==3.7.0 # via importlib-metadata diff --git a/plugins/flytekit-k8s-pod/requirements.txt b/plugins/flytekit-k8s-pod/requirements.txt index fa01bf43ca..6574f4ab66 100644 --- a/plugins/flytekit-k8s-pod/requirements.txt +++ b/plugins/flytekit-k8s-pod/requirements.txt @@ -6,11 +6,11 @@ # -e file:.#egg=flytekitplugins-pod # via -r requirements.in -arrow==1.2.1 +arrow==1.2.2 # via jinja2-time binaryornot==0.4.4 # via cookiecutter -cachetools==4.2.4 +cachetools==5.0.0 # via google-auth certifi==2021.10.8 # via @@ -20,7 +20,7 @@ cffi==1.15.0 # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.7 +charset-normalizer==2.0.11 # via requests checksumdir==1.2.0 # via flytekit @@ -32,33 +32,33 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.0.15 +croniter==1.2.0 # via flytekit -cryptography==35.0.0 +cryptography==36.0.1 # via secretstorage dataclasses-json==0.5.6 # via flytekit -decorator==5.1.0 +decorator==5.1.1 # via retry deprecated==1.2.13 # via flytekit -diskcache==5.2.1 +diskcache==5.4.0 # via flytekit docker-image-py==0.1.12 # via flytekit -docstring-parser==0.12 +docstring-parser==0.13 # via flytekit -flyteidl==0.21.8 +flyteidl==0.22.0 # via flytekit -flytekit==0.24.0 +flytekit==0.30.0 # via flytekitplugins-pod -google-auth==2.3.3 +google-auth==2.6.0 # via kubernetes -grpcio==1.41.1 +grpcio==1.43.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.8.2 +importlib-metadata==4.10.1 # via keyring jeepney==0.7.1 # via @@ -70,13 +70,13 @@ jinja2==3.0.3 # jinja2-time jinja2-time==0.2.0 # via cookiecutter -keyring==23.2.1 +keyring==23.5.0 # via flytekit -kubernetes==19.15.0 +kubernetes==21.7.0 # via flytekitplugins-pod markupsafe==2.0.1 # via jinja2 -marshmallow==3.14.0 +marshmallow==3.14.1 # via # dataclasses-json # marshmallow-enum @@ -87,25 +87,25 @@ marshmallow-jsonschema==0.13.0 # via flytekit mypy-extensions==0.4.3 # via typing-inspect -natsort==8.0.0 +natsort==8.1.0 # via flytekit -numpy==1.21.4 +numpy==1.22.2 # via # pandas # pyarrow -oauthlib==3.1.1 +oauthlib==3.2.0 # via requests-oauthlib -pandas==1.3.4 +pandas==1.4.0 # via flytekit poyo==0.5.0 # via cookiecutter -protobuf==3.19.1 +protobuf==3.19.4 # via # flyteidl # flytekit py==1.11.0 # via retry -pyarrow==6.0.0 +pyarrow==6.0.1 # via flytekit pyasn1==0.4.8 # via @@ -115,7 +115,7 @@ pyasn1-modules==0.2.8 # via google-auth pycparser==2.21 # via cffi -python-dateutil==2.8.1 +python-dateutil==2.8.2 # via # arrow # croniter @@ -128,65 +128,65 @@ python-slugify==5.0.2 # via cookiecutter pytimeparse==1.1.8 # via flytekit -pytz==2018.4 +pytz==2021.3 # via # flytekit # pandas pyyaml==6.0 # via kubernetes -regex==2021.11.10 +regex==2022.1.18 # via docker-image-py -requests==2.26.0 +requests==2.27.1 # via # cookiecutter # flytekit # kubernetes # requests-oauthlib # responses -requests-oauthlib==1.3.0 +requests-oauthlib==1.3.1 # via kubernetes -responses==0.15.0 +responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit -rsa==4.7.2 +rsa==4.8 # via google-auth secretstorage==3.3.1 # via keyring six==1.16.0 # via # cookiecutter - # flytekit # google-auth # grpcio # kubernetes # python-dateutil - # responses sortedcontainers==2.4.0 # via flytekit statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -typing-extensions==3.10.0.2 - # via typing-inspect +typing-extensions==4.0.1 + # via + # flytekit + # typing-inspect typing-inspect==0.7.1 # via dataclasses-json -urllib3==1.26.7 +urllib3==1.26.8 # via # flytekit # kubernetes # requests # responses -websocket-client==1.2.1 +websocket-client==1.2.3 # via kubernetes -wheel==0.37.0 +wheel==0.37.1 # via flytekit wrapt==1.13.3 # via # deprecated # flytekit -zipp==3.6.0 +zipp==3.7.0 # via importlib-metadata # The following packages are considered to be unsafe in a requirements file: diff --git a/plugins/flytekit-kf-mpi/requirements.txt b/plugins/flytekit-kf-mpi/requirements.txt index c65ce6c7b3..ef91243daf 100644 --- a/plugins/flytekit-kf-mpi/requirements.txt +++ b/plugins/flytekit-kf-mpi/requirements.txt @@ -6,7 +6,7 @@ # -e file:.#egg=flytekitplugins-kfmpi # via -r requirements.in -arrow==1.2.1 +arrow==1.2.2 # via jinja2-time binaryornot==0.4.4 # via cookiecutter @@ -16,7 +16,7 @@ cffi==1.15.0 # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.7 +charset-normalizer==2.0.11 # via requests checksumdir==1.2.0 # via flytekit @@ -28,33 +28,33 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.0.15 +croniter==1.2.0 # via flytekit -cryptography==35.0.0 +cryptography==36.0.1 # via secretstorage dataclasses-json==0.5.6 # via flytekit -decorator==5.1.0 +decorator==5.1.1 # via retry deprecated==1.2.13 # via flytekit -diskcache==5.2.1 +diskcache==5.4.0 # via flytekit docker-image-py==0.1.12 # via flytekit -docstring-parser==0.12 +docstring-parser==0.13 # via flytekit -flyteidl==0.21.8 +flyteidl==0.22.0 # via # flytekit # flytekitplugins-kfmpi -flytekit==0.24.0 +flytekit==0.30.0 # via flytekitplugins-kfmpi -grpcio==1.41.1 +grpcio==1.43.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.8.2 +importlib-metadata==4.10.1 # via keyring jeepney==0.7.1 # via @@ -66,11 +66,11 @@ jinja2==3.0.3 # jinja2-time jinja2-time==0.2.0 # via cookiecutter -keyring==23.2.1 +keyring==23.5.0 # via flytekit markupsafe==2.0.1 # via jinja2 -marshmallow==3.14.0 +marshmallow==3.14.1 # via # dataclasses-json # marshmallow-enum @@ -81,27 +81,27 @@ marshmallow-jsonschema==0.13.0 # via flytekit mypy-extensions==0.4.3 # via typing-inspect -natsort==8.0.0 +natsort==8.1.0 # via flytekit -numpy==1.21.4 +numpy==1.22.2 # via # pandas # pyarrow -pandas==1.3.4 +pandas==1.4.0 # via flytekit poyo==0.5.0 # via cookiecutter -protobuf==3.19.1 +protobuf==3.19.4 # via # flyteidl # flytekit py==1.11.0 # via retry -pyarrow==6.0.0 +pyarrow==6.0.1 # via flytekit pycparser==2.21 # via cffi -python-dateutil==2.8.1 +python-dateutil==2.8.2 # via # arrow # croniter @@ -113,18 +113,18 @@ python-slugify==5.0.2 # via cookiecutter pytimeparse==1.1.8 # via flytekit -pytz==2018.4 +pytz==2021.3 # via # flytekit # pandas -regex==2021.11.10 +regex==2022.1.18 # via docker-image-py -requests==2.26.0 +requests==2.27.1 # via # cookiecutter # flytekit # responses -responses==0.15.0 +responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit @@ -133,30 +133,30 @@ secretstorage==3.3.1 six==1.16.0 # via # cookiecutter - # flytekit # grpcio # python-dateutil - # responses sortedcontainers==2.4.0 # via flytekit statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -typing-extensions==3.10.0.2 - # via typing-inspect +typing-extensions==4.0.1 + # via + # flytekit + # typing-inspect typing-inspect==0.7.1 # via dataclasses-json -urllib3==1.26.7 +urllib3==1.26.8 # via # flytekit # requests # responses -wheel==0.37.0 +wheel==0.37.1 # via flytekit wrapt==1.13.3 # via # deprecated # flytekit -zipp==3.6.0 +zipp==3.7.0 # via importlib-metadata diff --git a/plugins/flytekit-kf-pytorch/requirements.txt b/plugins/flytekit-kf-pytorch/requirements.txt index 2cdcd130c9..9fc8665f43 100644 --- a/plugins/flytekit-kf-pytorch/requirements.txt +++ b/plugins/flytekit-kf-pytorch/requirements.txt @@ -6,7 +6,7 @@ # -e file:.#egg=flytekitplugins-kfpytorch # via -r requirements.in -arrow==1.2.1 +arrow==1.2.2 # via jinja2-time binaryornot==0.4.4 # via cookiecutter @@ -16,7 +16,7 @@ cffi==1.15.0 # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.7 +charset-normalizer==2.0.11 # via requests checksumdir==1.2.0 # via flytekit @@ -28,31 +28,31 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.0.15 +croniter==1.2.0 # via flytekit -cryptography==35.0.0 +cryptography==36.0.1 # via secretstorage dataclasses-json==0.5.6 # via flytekit -decorator==5.1.0 +decorator==5.1.1 # via retry deprecated==1.2.13 # via flytekit -diskcache==5.2.1 +diskcache==5.4.0 # via flytekit docker-image-py==0.1.12 # via flytekit -docstring-parser==0.12 +docstring-parser==0.13 # via flytekit -flyteidl==0.21.8 +flyteidl==0.22.0 # via flytekit -flytekit==0.24.0 +flytekit==0.30.0 # via flytekitplugins-kfpytorch -grpcio==1.41.1 +grpcio==1.43.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.8.2 +importlib-metadata==4.10.1 # via keyring jeepney==0.7.1 # via @@ -64,11 +64,11 @@ jinja2==3.0.3 # jinja2-time jinja2-time==0.2.0 # via cookiecutter -keyring==23.2.1 +keyring==23.5.0 # via flytekit markupsafe==2.0.1 # via jinja2 -marshmallow==3.14.0 +marshmallow==3.14.1 # via # dataclasses-json # marshmallow-enum @@ -79,27 +79,27 @@ marshmallow-jsonschema==0.13.0 # via flytekit mypy-extensions==0.4.3 # via typing-inspect -natsort==8.0.0 +natsort==8.1.0 # via flytekit -numpy==1.21.4 +numpy==1.22.2 # via # pandas # pyarrow -pandas==1.3.4 +pandas==1.4.0 # via flytekit poyo==0.5.0 # via cookiecutter -protobuf==3.19.1 +protobuf==3.19.4 # via # flyteidl # flytekit py==1.11.0 # via retry -pyarrow==6.0.0 +pyarrow==6.0.1 # via flytekit pycparser==2.21 # via cffi -python-dateutil==2.8.1 +python-dateutil==2.8.2 # via # arrow # croniter @@ -111,18 +111,18 @@ python-slugify==5.0.2 # via cookiecutter pytimeparse==1.1.8 # via flytekit -pytz==2018.4 +pytz==2021.3 # via # flytekit # pandas -regex==2021.11.10 +regex==2022.1.18 # via docker-image-py -requests==2.26.0 +requests==2.27.1 # via # cookiecutter # flytekit # responses -responses==0.15.0 +responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit @@ -131,30 +131,30 @@ secretstorage==3.3.1 six==1.16.0 # via # cookiecutter - # flytekit # grpcio # python-dateutil - # responses sortedcontainers==2.4.0 # via flytekit statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -typing-extensions==3.10.0.2 - # via typing-inspect +typing-extensions==4.0.1 + # via + # flytekit + # typing-inspect typing-inspect==0.7.1 # via dataclasses-json -urllib3==1.26.7 +urllib3==1.26.8 # via # flytekit # requests # responses -wheel==0.37.0 +wheel==0.37.1 # via flytekit wrapt==1.13.3 # via # deprecated # flytekit -zipp==3.6.0 +zipp==3.7.0 # via importlib-metadata diff --git a/plugins/flytekit-kf-tensorflow/requirements.txt b/plugins/flytekit-kf-tensorflow/requirements.txt index ac3a64ed65..ae428e9215 100644 --- a/plugins/flytekit-kf-tensorflow/requirements.txt +++ b/plugins/flytekit-kf-tensorflow/requirements.txt @@ -6,7 +6,7 @@ # -e file:.#egg=flytekitplugins-kftensorflow # via -r requirements.in -arrow==1.2.1 +arrow==1.2.2 # via jinja2-time binaryornot==0.4.4 # via cookiecutter @@ -16,7 +16,7 @@ cffi==1.15.0 # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.7 +charset-normalizer==2.0.11 # via requests checksumdir==1.2.0 # via flytekit @@ -28,31 +28,31 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.0.15 +croniter==1.2.0 # via flytekit -cryptography==35.0.0 +cryptography==36.0.1 # via secretstorage dataclasses-json==0.5.6 # via flytekit -decorator==5.1.0 +decorator==5.1.1 # via retry deprecated==1.2.13 # via flytekit -diskcache==5.2.1 +diskcache==5.4.0 # via flytekit docker-image-py==0.1.12 # via flytekit -docstring-parser==0.12 +docstring-parser==0.13 # via flytekit -flyteidl==0.21.8 +flyteidl==0.22.0 # via flytekit -flytekit==0.24.0 +flytekit==0.30.0 # via flytekitplugins-kftensorflow -grpcio==1.41.1 +grpcio==1.43.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.8.2 +importlib-metadata==4.10.1 # via keyring jeepney==0.7.1 # via @@ -64,11 +64,11 @@ jinja2==3.0.3 # jinja2-time jinja2-time==0.2.0 # via cookiecutter -keyring==23.2.1 +keyring==23.5.0 # via flytekit markupsafe==2.0.1 # via jinja2 -marshmallow==3.14.0 +marshmallow==3.14.1 # via # dataclasses-json # marshmallow-enum @@ -79,27 +79,27 @@ marshmallow-jsonschema==0.13.0 # via flytekit mypy-extensions==0.4.3 # via typing-inspect -natsort==8.0.0 +natsort==8.1.0 # via flytekit -numpy==1.21.4 +numpy==1.22.2 # via # pandas # pyarrow -pandas==1.3.4 +pandas==1.4.0 # via flytekit poyo==0.5.0 # via cookiecutter -protobuf==3.19.1 +protobuf==3.19.4 # via # flyteidl # flytekit py==1.11.0 # via retry -pyarrow==6.0.0 +pyarrow==6.0.1 # via flytekit pycparser==2.21 # via cffi -python-dateutil==2.8.1 +python-dateutil==2.8.2 # via # arrow # croniter @@ -111,18 +111,18 @@ python-slugify==5.0.2 # via cookiecutter pytimeparse==1.1.8 # via flytekit -pytz==2018.4 +pytz==2021.3 # via # flytekit # pandas -regex==2021.11.10 +regex==2022.1.18 # via docker-image-py -requests==2.26.0 +requests==2.27.1 # via # cookiecutter # flytekit # responses -responses==0.15.0 +responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit @@ -131,30 +131,30 @@ secretstorage==3.3.1 six==1.16.0 # via # cookiecutter - # flytekit # grpcio # python-dateutil - # responses sortedcontainers==2.4.0 # via flytekit statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -typing-extensions==3.10.0.2 - # via typing-inspect +typing-extensions==4.0.1 + # via + # flytekit + # typing-inspect typing-inspect==0.7.1 # via dataclasses-json -urllib3==1.26.7 +urllib3==1.26.8 # via # flytekit # requests # responses -wheel==0.37.0 +wheel==0.37.1 # via flytekit wrapt==1.13.3 # via # deprecated # flytekit -zipp==3.6.0 +zipp==3.7.0 # via importlib-metadata diff --git a/plugins/flytekit-modin/requirements.txt b/plugins/flytekit-modin/requirements.txt index ff80b3a0f6..adaa0b7307 100644 --- a/plugins/flytekit-modin/requirements.txt +++ b/plugins/flytekit-modin/requirements.txt @@ -6,9 +6,9 @@ # -e file:.#egg=flytekitplugins-modin # via -r requirements.in -arrow==1.2.1 +arrow==1.2.2 # via jinja2-time -attrs==21.2.0 +attrs==21.4.0 # via # jsonschema # ray @@ -20,7 +20,7 @@ cffi==1.15.0 # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.7 +charset-normalizer==2.0.11 # via requests checksumdir==1.2.0 # via flytekit @@ -33,37 +33,41 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.0.15 +croniter==1.2.0 # via flytekit -cryptography==35.0.0 +cryptography==36.0.1 # via secretstorage dataclasses-json==0.5.6 # via flytekit -decorator==5.1.0 +decorator==5.1.1 # via retry deprecated==1.2.13 - # via flytekit -diskcache==5.2.1 + # via + # flytekit + # redis +diskcache==5.4.0 # via flytekit docker-image-py==0.1.12 # via flytekit -docstring-parser==0.12 +docstring-parser==0.13 # via flytekit -filelock==3.3.2 +filelock==3.4.2 # via ray -flyteidl==0.21.8 +flyteidl==0.22.0 # via flytekit -flytekit==0.24.0 - # via flytekitplugins-modin -fsspec==2021.11.0 +flytekit==0.30.0 # via flytekitplugins-modin -grpcio==1.41.1 +fsspec==2022.1.0 + # via + # flytekitplugins-modin + # modin +grpcio==1.43.0 # via # flytekit # ray idna==3.3 # via requests -importlib-metadata==4.8.2 +importlib-metadata==4.10.1 # via keyring jeepney==0.7.1 # via @@ -75,13 +79,13 @@ jinja2==3.0.3 # jinja2-time jinja2-time==0.2.0 # via cookiecutter -jsonschema==4.2.1 +jsonschema==4.4.0 # via ray -keyring==23.2.1 +keyring==23.5.0 # via flytekit markupsafe==2.0.1 # via jinja2 -marshmallow==3.14.0 +marshmallow==3.14.1 # via # dataclasses-json # marshmallow-enum @@ -90,44 +94,46 @@ marshmallow-enum==1.5.1 # via dataclasses-json marshmallow-jsonschema==0.13.0 # via flytekit -modin==0.11.3 +modin==0.13.1 # via flytekitplugins-modin -msgpack==1.0.2 +msgpack==1.0.3 # via ray mypy-extensions==0.4.3 # via typing-inspect -natsort==8.0.0 +natsort==8.1.0 # via flytekit -numpy==1.21.4 +numpy==1.22.2 # via # modin # pandas # pyarrow # ray -packaging==21.2 - # via modin -pandas==1.3.4 +packaging==21.3 + # via + # modin + # redis +pandas==1.4.0 # via # flytekit # modin poyo==0.5.0 # via cookiecutter -protobuf==3.19.1 +protobuf==3.19.4 # via # flyteidl # flytekit # ray py==1.11.0 # via retry -pyarrow==6.0.0 +pyarrow==6.0.1 # via flytekit pycparser==2.21 # via cffi -pyparsing==2.4.7 +pyparsing==3.0.7 # via packaging -pyrsistent==0.18.0 +pyrsistent==0.18.1 # via jsonschema -python-dateutil==2.8.1 +python-dateutil==2.8.2 # via # arrow # croniter @@ -139,24 +145,24 @@ python-slugify==5.0.2 # via cookiecutter pytimeparse==1.1.8 # via flytekit -pytz==2018.4 +pytz==2021.3 # via # flytekit # pandas pyyaml==6.0 # via ray -ray==1.8.0 +ray==1.10.0 # via flytekitplugins-modin -redis==3.5.3 +redis==4.1.2 # via ray -regex==2021.11.10 +regex==2022.1.18 # via docker-image-py -requests==2.26.0 +requests==2.27.1 # via # cookiecutter # flytekit # responses -responses==0.15.0 +responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit @@ -165,30 +171,30 @@ secretstorage==3.3.1 six==1.16.0 # via # cookiecutter - # flytekit # grpcio # python-dateutil - # responses sortedcontainers==2.4.0 # via flytekit statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -typing-extensions==3.10.0.2 - # via typing-inspect +typing-extensions==4.0.1 + # via + # flytekit + # typing-inspect typing-inspect==0.7.1 # via dataclasses-json -urllib3==1.26.7 +urllib3==1.26.8 # via # flytekit # requests # responses -wheel==0.37.0 +wheel==0.37.1 # via flytekit wrapt==1.13.3 # via # deprecated # flytekit -zipp==3.6.0 +zipp==3.7.0 # via importlib-metadata diff --git a/plugins/flytekit-pandera/requirements.txt b/plugins/flytekit-pandera/requirements.txt index de61384e56..82df51d30e 100644 --- a/plugins/flytekit-pandera/requirements.txt +++ b/plugins/flytekit-pandera/requirements.txt @@ -6,7 +6,7 @@ # -e file:.#egg=flytekitplugins-pandera # via -r requirements.in -arrow==1.2.1 +arrow==1.2.2 # via jinja2-time binaryornot==0.4.4 # via cookiecutter @@ -16,7 +16,7 @@ cffi==1.15.0 # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.7 +charset-normalizer==2.0.11 # via requests checksumdir==1.2.0 # via flytekit @@ -28,31 +28,31 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.0.15 +croniter==1.2.0 # via flytekit -cryptography==35.0.0 +cryptography==36.0.1 # via secretstorage dataclasses-json==0.5.6 # via flytekit -decorator==5.1.0 +decorator==5.1.1 # via retry deprecated==1.2.13 # via flytekit -diskcache==5.2.1 +diskcache==5.4.0 # via flytekit docker-image-py==0.1.12 # via flytekit -docstring-parser==0.12 +docstring-parser==0.13 # via flytekit -flyteidl==0.21.8 +flyteidl==0.22.0 # via flytekit -flytekit==0.24.0 +flytekit==0.30.0 # via flytekitplugins-pandera -grpcio==1.41.1 +grpcio==1.43.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.8.2 +importlib-metadata==4.10.1 # via keyring jeepney==0.7.1 # via @@ -64,11 +64,11 @@ jinja2==3.0.3 # jinja2-time jinja2-time==0.2.0 # via cookiecutter -keyring==23.2.1 +keyring==23.5.0 # via flytekit markupsafe==2.0.1 # via jinja2 -marshmallow==3.14.0 +marshmallow==3.14.1 # via # dataclasses-json # marshmallow-enum @@ -79,38 +79,38 @@ marshmallow-jsonschema==0.13.0 # via flytekit mypy-extensions==0.4.3 # via typing-inspect -natsort==8.0.0 +natsort==8.1.0 # via flytekit -numpy==1.21.4 +numpy==1.22.2 # via # pandas # pandera # pyarrow -packaging==21.2 +packaging==21.3 # via pandera -pandas==1.3.4 +pandas==1.4.0 # via # flytekit # pandera -pandera==0.7.2 +pandera==0.8.1 # via flytekitplugins-pandera poyo==0.5.0 # via cookiecutter -protobuf==3.19.1 +protobuf==3.19.4 # via # flyteidl # flytekit py==1.11.0 # via retry -pyarrow==6.0.0 +pyarrow==6.0.1 # via # flytekit # pandera pycparser==2.21 # via cffi -pyparsing==2.4.7 +pyparsing==3.0.7 # via packaging -python-dateutil==2.8.1 +python-dateutil==2.8.2 # via # arrow # croniter @@ -122,18 +122,18 @@ python-slugify==5.0.2 # via cookiecutter pytimeparse==1.1.8 # via flytekit -pytz==2018.4 +pytz==2021.3 # via # flytekit # pandas -regex==2021.11.10 +regex==2022.1.18 # via docker-image-py -requests==2.26.0 +requests==2.27.1 # via # cookiecutter # flytekit # responses -responses==0.15.0 +responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit @@ -142,33 +142,33 @@ secretstorage==3.3.1 six==1.16.0 # via # cookiecutter - # flytekit # grpcio # python-dateutil - # responses sortedcontainers==2.4.0 # via flytekit statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -typing-extensions==3.10.0.2 - # via typing-inspect +typing-extensions==4.0.1 + # via + # flytekit + # typing-inspect typing-inspect==0.7.1 # via # dataclasses-json # pandera -urllib3==1.26.7 +urllib3==1.26.8 # via # flytekit # requests # responses -wheel==0.37.0 +wheel==0.37.1 # via flytekit wrapt==1.13.3 # via # deprecated # flytekit # pandera -zipp==3.6.0 +zipp==3.7.0 # via importlib-metadata diff --git a/plugins/flytekit-papermill/requirements.txt b/plugins/flytekit-papermill/requirements.txt index a4b8b61e3d..b06adde647 100644 --- a/plugins/flytekit-papermill/requirements.txt +++ b/plugins/flytekit-papermill/requirements.txt @@ -8,27 +8,27 @@ # via -r requirements.in ansiwrap==0.8.4 # via papermill -appnope==0.1.2 - # via - # ipykernel - # ipython -arrow==1.2.1 +arrow==1.2.2 # via jinja2-time -attrs==21.2.0 +asttokens==2.0.5 + # via stack-data +attrs==21.4.0 # via jsonschema backcall==0.2.0 # via ipython binaryornot==0.4.4 # via cookiecutter -black==21.10b0 - # via papermill +black==21.12b0 + # via ipython bleach==4.1.0 # via nbconvert certifi==2021.10.8 # via requests +cffi==1.15.0 + # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.7 +charset-normalizer==2.0.10 # via requests checksumdir==1.2.0 # via flytekit @@ -42,13 +42,15 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.0.15 +croniter==1.2.0 # via flytekit +cryptography==36.0.1 + # via secretstorage dataclasses-json==0.5.6 # via flytekit debugpy==1.5.1 # via ipykernel -decorator==5.1.0 +decorator==5.1.1 # via # ipython # retry @@ -56,35 +58,41 @@ defusedxml==0.7.1 # via nbconvert deprecated==1.2.13 # via flytekit -diskcache==5.2.1 +diskcache==5.4.0 # via flytekit docker-image-py==0.1.12 # via flytekit -docstring-parser==0.12 +docstring-parser==0.13 # via flytekit entrypoints==0.3 # via # jupyter-client # nbconvert # papermill -flyteidl==0.21.23 +executing==0.8.2 + # via stack-data +flyteidl==0.21.24 # via flytekit -flytekit==0.24.0 +flytekit==0.26.1 # via flytekitplugins-papermill -grpcio==1.41.1 +grpcio==1.43.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.8.2 +importlib-metadata==4.10.1 # via keyring -ipykernel==6.5.0 +ipykernel==6.7.0 # via flytekitplugins-papermill -ipython==7.31.1 +ipython==8.0.1 # via ipykernel ipython-genutils==0.2.0 # via nbformat -jedi==0.18.0 +jedi==0.18.1 # via ipython +jeepney==0.7.1 + # via + # keyring + # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -92,9 +100,9 @@ jinja2==3.0.3 # nbconvert jinja2-time==0.2.0 # via cookiecutter -jsonschema==4.2.1 +jsonschema==4.4.0 # via nbformat -jupyter-client==7.0.6 +jupyter-client==7.1.2 # via # ipykernel # nbclient @@ -105,11 +113,11 @@ jupyter-core==4.9.1 # nbformat jupyterlab-pygments==0.1.2 # via nbconvert -keyring==23.2.1 +keyring==23.5.0 # via flytekit markupsafe==2.0.1 # via jinja2 -marshmallow==3.14.0 +marshmallow==3.14.1 # via # dataclasses-json # marshmallow-enum @@ -128,36 +136,37 @@ mypy-extensions==0.4.3 # via # black # typing-inspect -natsort==8.0.0 +natsort==8.0.2 # via flytekit -nbclient==0.5.5 +nbclient==0.5.10 # via # nbconvert # papermill -nbconvert==6.2.0 +nbconvert==6.4.1 # via flytekitplugins-papermill nbformat==5.1.3 # via # nbclient # nbconvert # papermill -nest-asyncio==1.5.1 +nest-asyncio==1.5.4 # via + # ipykernel # jupyter-client # nbclient -numpy==1.21.4 +numpy==1.22.1 # via # pandas # pyarrow -packaging==21.2 +packaging==21.3 # via bleach -pandas==1.3.4 +pandas==1.4.0 # via flytekit pandocfilters==1.5.0 # via nbconvert -papermill==2.3.3 +papermill==2.3.4 # via flytekitplugins-papermill -parso==0.8.2 +parso==0.8.3 # via jedi pathspec==0.9.0 # via black @@ -165,30 +174,34 @@ pexpect==4.8.0 # via ipython pickleshare==0.7.5 # via ipython -platformdirs==2.4.0 +platformdirs==2.4.1 # via black poyo==0.5.0 # via cookiecutter -prompt-toolkit==3.0.22 +prompt-toolkit==3.0.26 # via ipython -protobuf==3.19.1 +protobuf==3.19.4 # via # flyteidl # flytekit ptyprocess==0.7.0 # via pexpect +pure-eval==0.2.2 + # via stack-data py==1.11.0 # via retry -pyarrow==6.0.0 +pyarrow==6.0.1 # via flytekit -pygments==2.10.0 +pycparser==2.21 + # via cffi +pygments==2.11.2 # via # ipython # jupyterlab-pygments # nbconvert -pyparsing==2.4.7 +pyparsing==3.0.7 # via packaging -pyrsistent==0.18.0 +pyrsistent==0.18.1 # via jsonschema python-dateutil==2.8.1 # via @@ -203,7 +216,7 @@ python-slugify==5.0.2 # via cookiecutter pytimeparse==1.1.8 # via flytekit -pytz==2018.4 +pytz==2021.3 # via # flytekit # pandas @@ -211,22 +224,23 @@ pyyaml==6.0 # via papermill pyzmq==22.3.0 # via jupyter-client -regex==2021.11.10 - # via - # black - # docker-image-py -requests==2.26.0 +regex==2022.1.18 + # via docker-image-py +requests==2.27.1 # via # cookiecutter # flytekit # papermill # responses -responses==0.15.0 +responses==0.17.0 # via flytekit retry==0.9.2 # via flytekit +secretstorage==3.3.1 + # via keyring six==1.16.0 # via + # asttokens # bleach # cookiecutter # flytekit @@ -235,6 +249,8 @@ six==1.16.0 # responses sortedcontainers==2.4.0 # via flytekit +stack-data==0.1.4 + # via ipython statsd==3.3.0 # via flytekit tenacity==8.0.1 @@ -245,7 +261,7 @@ text-unidecode==1.3 # via python-slugify textwrap3==0.9.2 # via ansiwrap -tomli==1.2.2 +tomli==1.2.3 # via black tornado==6.1 # via @@ -263,13 +279,13 @@ traitlets==5.1.1 # nbclient # nbconvert # nbformat -typing-extensions==3.10.0.2 +typing-extensions==4.0.1 # via # black # typing-inspect typing-inspect==0.7.1 # via dataclasses-json -urllib3==1.26.7 +urllib3==1.26.8 # via # flytekit # requests @@ -278,13 +294,13 @@ wcwidth==0.2.5 # via prompt-toolkit webencodings==0.5.1 # via bleach -wheel==0.37.0 +wheel==0.37.1 # via flytekit wrapt==1.13.3 # via # deprecated # flytekit -zipp==3.6.0 +zipp==3.7.0 # via importlib-metadata # The following packages are considered to be unsafe in a requirements file: diff --git a/plugins/flytekit-snowflake/requirements.txt b/plugins/flytekit-snowflake/requirements.txt index 04a2dfec1d..b563e54814 100644 --- a/plugins/flytekit-snowflake/requirements.txt +++ b/plugins/flytekit-snowflake/requirements.txt @@ -6,7 +6,7 @@ # -e file:.#egg=flytekitplugins-snowflake # via -r requirements.in -arrow==1.2.1 +arrow==1.2.2 # via jinja2-time binaryornot==0.4.4 # via cookiecutter @@ -16,7 +16,7 @@ cffi==1.15.0 # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.7 +charset-normalizer==2.0.11 # via requests checksumdir==1.2.0 # via flytekit @@ -28,31 +28,31 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.0.15 +croniter==1.2.0 # via flytekit -cryptography==35.0.0 +cryptography==36.0.1 # via secretstorage dataclasses-json==0.5.6 # via flytekit -decorator==5.1.0 +decorator==5.1.1 # via retry deprecated==1.2.13 # via flytekit -diskcache==5.2.1 +diskcache==5.4.0 # via flytekit docker-image-py==0.1.12 # via flytekit -docstring-parser==0.12 +docstring-parser==0.13 # via flytekit -flyteidl==0.21.8 +flyteidl==0.22.0 # via flytekit -flytekit==0.24.0 +flytekit==0.30.0 # via flytekitplugins-snowflake -grpcio==1.41.1 +grpcio==1.43.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.8.2 +importlib-metadata==4.10.1 # via keyring jeepney==0.7.1 # via @@ -64,11 +64,11 @@ jinja2==3.0.3 # jinja2-time jinja2-time==0.2.0 # via cookiecutter -keyring==23.2.1 +keyring==23.5.0 # via flytekit markupsafe==2.0.1 # via jinja2 -marshmallow==3.14.0 +marshmallow==3.14.1 # via # dataclasses-json # marshmallow-enum @@ -79,27 +79,27 @@ marshmallow-jsonschema==0.13.0 # via flytekit mypy-extensions==0.4.3 # via typing-inspect -natsort==8.0.0 +natsort==8.1.0 # via flytekit -numpy==1.21.4 +numpy==1.22.2 # via # pandas # pyarrow -pandas==1.3.4 +pandas==1.4.0 # via flytekit poyo==0.5.0 # via cookiecutter -protobuf==3.19.1 +protobuf==3.19.4 # via # flyteidl # flytekit py==1.11.0 # via retry -pyarrow==6.0.0 +pyarrow==6.0.1 # via flytekit pycparser==2.21 # via cffi -python-dateutil==2.8.1 +python-dateutil==2.8.2 # via # arrow # croniter @@ -111,18 +111,18 @@ python-slugify==5.0.2 # via cookiecutter pytimeparse==1.1.8 # via flytekit -pytz==2018.4 +pytz==2021.3 # via # flytekit # pandas -regex==2021.11.10 +regex==2022.1.18 # via docker-image-py -requests==2.26.0 +requests==2.27.1 # via # cookiecutter # flytekit # responses -responses==0.15.0 +responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit @@ -131,30 +131,30 @@ secretstorage==3.3.1 six==1.16.0 # via # cookiecutter - # flytekit # grpcio # python-dateutil - # responses sortedcontainers==2.4.0 # via flytekit statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -typing-extensions==3.10.0.2 - # via typing-inspect +typing-extensions==4.0.1 + # via + # flytekit + # typing-inspect typing-inspect==0.7.1 # via dataclasses-json -urllib3==1.26.7 +urllib3==1.26.8 # via # flytekit # requests # responses -wheel==0.37.0 +wheel==0.37.1 # via flytekit wrapt==1.13.3 # via # deprecated # flytekit -zipp==3.6.0 +zipp==3.7.0 # via importlib-metadata diff --git a/plugins/flytekit-spark/requirements.txt b/plugins/flytekit-spark/requirements.txt index df2baaa778..fed298d4b4 100644 --- a/plugins/flytekit-spark/requirements.txt +++ b/plugins/flytekit-spark/requirements.txt @@ -6,7 +6,7 @@ # -e file:.#egg=flytekitplugins-spark # via -r requirements.in -arrow==1.2.1 +arrow==1.2.2 # via jinja2-time binaryornot==0.4.4 # via cookiecutter @@ -16,7 +16,7 @@ cffi==1.15.0 # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.7 +charset-normalizer==2.0.11 # via requests checksumdir==1.2.0 # via flytekit @@ -28,31 +28,31 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.0.15 +croniter==1.2.0 # via flytekit -cryptography==35.0.0 +cryptography==36.0.1 # via secretstorage dataclasses-json==0.5.6 # via flytekit -decorator==5.1.0 +decorator==5.1.1 # via retry deprecated==1.2.13 # via flytekit -diskcache==5.2.1 +diskcache==5.4.0 # via flytekit docker-image-py==0.1.12 # via flytekit -docstring-parser==0.12 +docstring-parser==0.13 # via flytekit -flyteidl==0.21.8 +flyteidl==0.22.0 # via flytekit -flytekit==0.24.0 +flytekit==0.30.0 # via flytekitplugins-spark -grpcio==1.41.1 +grpcio==1.43.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.8.2 +importlib-metadata==4.10.1 # via keyring jeepney==0.7.1 # via @@ -64,11 +64,11 @@ jinja2==3.0.3 # jinja2-time jinja2-time==0.2.0 # via cookiecutter -keyring==23.2.1 +keyring==23.5.0 # via flytekit markupsafe==2.0.1 # via jinja2 -marshmallow==3.14.0 +marshmallow==3.14.1 # via # dataclasses-json # marshmallow-enum @@ -79,31 +79,31 @@ marshmallow-jsonschema==0.13.0 # via flytekit mypy-extensions==0.4.3 # via typing-inspect -natsort==8.0.0 +natsort==8.1.0 # via flytekit -numpy==1.21.4 +numpy==1.22.2 # via # pandas # pyarrow -pandas==1.3.4 +pandas==1.4.0 # via flytekit poyo==0.5.0 # via cookiecutter -protobuf==3.19.1 +protobuf==3.19.4 # via # flyteidl # flytekit py==1.11.0 # via retry -py4j==0.10.9.2 +py4j==0.10.9.3 # via pyspark -pyarrow==6.0.0 +pyarrow==6.0.1 # via flytekit pycparser==2.21 # via cffi -pyspark==3.2.0 +pyspark==3.2.1 # via flytekitplugins-spark -python-dateutil==2.8.1 +python-dateutil==2.8.2 # via # arrow # croniter @@ -115,18 +115,18 @@ python-slugify==5.0.2 # via cookiecutter pytimeparse==1.1.8 # via flytekit -pytz==2018.4 +pytz==2021.3 # via # flytekit # pandas -regex==2021.11.10 +regex==2022.1.18 # via docker-image-py -requests==2.26.0 +requests==2.27.1 # via # cookiecutter # flytekit # responses -responses==0.15.0 +responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit @@ -135,30 +135,30 @@ secretstorage==3.3.1 six==1.16.0 # via # cookiecutter - # flytekit # grpcio # python-dateutil - # responses sortedcontainers==2.4.0 # via flytekit statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -typing-extensions==3.10.0.2 - # via typing-inspect +typing-extensions==4.0.1 + # via + # flytekit + # typing-inspect typing-inspect==0.7.1 # via dataclasses-json -urllib3==1.26.7 +urllib3==1.26.8 # via # flytekit # requests # responses -wheel==0.37.0 +wheel==0.37.1 # via flytekit wrapt==1.13.3 # via # deprecated # flytekit -zipp==3.6.0 +zipp==3.7.0 # via importlib-metadata diff --git a/plugins/flytekit-sqlalchemy/requirements.txt b/plugins/flytekit-sqlalchemy/requirements.txt index bf4f5f0590..134f8960e8 100644 --- a/plugins/flytekit-sqlalchemy/requirements.txt +++ b/plugins/flytekit-sqlalchemy/requirements.txt @@ -6,7 +6,7 @@ # -e file:.#egg=flytekitplugins-sqlalchemy # via -r requirements.in -arrow==1.2.1 +arrow==1.2.2 # via jinja2-time binaryornot==0.4.4 # via cookiecutter @@ -16,7 +16,7 @@ cffi==1.15.0 # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.7 +charset-normalizer==2.0.11 # via requests checksumdir==1.2.0 # via flytekit @@ -28,33 +28,33 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.0.15 +croniter==1.2.0 # via flytekit -cryptography==35.0.0 +cryptography==36.0.1 # via secretstorage dataclasses-json==0.5.6 # via flytekit -decorator==5.1.0 +decorator==5.1.1 # via retry deprecated==1.2.13 # via flytekit -diskcache==5.2.1 +diskcache==5.4.0 # via flytekit docker-image-py==0.1.12 # via flytekit -docstring-parser==0.12 +docstring-parser==0.13 # via flytekit -flyteidl==0.21.8 +flyteidl==0.22.0 # via flytekit -flytekit==0.24.0 +flytekit==0.30.0 # via flytekitplugins-sqlalchemy greenlet==1.1.2 # via sqlalchemy -grpcio==1.41.1 +grpcio==1.43.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.8.2 +importlib-metadata==4.10.1 # via keyring jeepney==0.7.1 # via @@ -66,11 +66,11 @@ jinja2==3.0.3 # jinja2-time jinja2-time==0.2.0 # via cookiecutter -keyring==23.2.1 +keyring==23.5.0 # via flytekit markupsafe==2.0.1 # via jinja2 -marshmallow==3.14.0 +marshmallow==3.14.1 # via # dataclasses-json # marshmallow-enum @@ -81,27 +81,27 @@ marshmallow-jsonschema==0.13.0 # via flytekit mypy-extensions==0.4.3 # via typing-inspect -natsort==8.0.0 +natsort==8.1.0 # via flytekit -numpy==1.21.4 +numpy==1.22.2 # via # pandas # pyarrow -pandas==1.3.4 +pandas==1.4.0 # via flytekit poyo==0.5.0 # via cookiecutter -protobuf==3.19.1 +protobuf==3.19.4 # via # flyteidl # flytekit py==1.11.0 # via retry -pyarrow==6.0.0 +pyarrow==6.0.1 # via flytekit pycparser==2.21 # via cffi -python-dateutil==2.8.1 +python-dateutil==2.8.2 # via # arrow # croniter @@ -113,18 +113,18 @@ python-slugify==5.0.2 # via cookiecutter pytimeparse==1.1.8 # via flytekit -pytz==2018.4 +pytz==2021.3 # via # flytekit # pandas -regex==2021.11.10 +regex==2022.1.18 # via docker-image-py -requests==2.26.0 +requests==2.27.1 # via # cookiecutter # flytekit # responses -responses==0.15.0 +responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit @@ -133,32 +133,32 @@ secretstorage==3.3.1 six==1.16.0 # via # cookiecutter - # flytekit # grpcio # python-dateutil - # responses sortedcontainers==2.4.0 # via flytekit -sqlalchemy==1.4.26 +sqlalchemy==1.4.31 # via flytekitplugins-sqlalchemy statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -typing-extensions==3.10.0.2 - # via typing-inspect +typing-extensions==4.0.1 + # via + # flytekit + # typing-inspect typing-inspect==0.7.1 # via dataclasses-json -urllib3==1.26.7 +urllib3==1.26.8 # via # flytekit # requests # responses -wheel==0.37.0 +wheel==0.37.1 # via flytekit wrapt==1.13.3 # via # deprecated # flytekit -zipp==3.6.0 +zipp==3.7.0 # via importlib-metadata diff --git a/requirements-spark2.txt b/requirements-spark2.txt index bdef3764f0..3e15827223 100644 --- a/requirements-spark2.txt +++ b/requirements-spark2.txt @@ -8,107 +8,56 @@ # via # -r requirements-spark2.in # -r requirements.in -ansiwrap==0.8.4 - # via papermill -arrow==1.2.1 +arrow==1.2.2 # via jinja2-time attrs==20.3.0 # via # -r requirements.in # jsonschema -backcall==0.2.0 - # via ipython -bcrypt==3.2.0 - # via paramiko binaryornot==0.4.4 # via cookiecutter -black==21.12b0 - # via papermill -bleach==4.1.0 - # via nbconvert -boto3==1.20.26 - # via sagemaker-training -botocore==1.23.26 - # via - # boto3 - # s3transfer certifi==2021.10.8 # via requests cffi==1.15.0 - # via - # bcrypt - # cryptography - # pynacl + # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.9 +charset-normalizer==2.0.11 # via requests checksumdir==1.2.0 # via flytekit -click==7.1.2 +click==8.0.3 # via - # black # cookiecutter # flytekit - # hmsclient - # papermill cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.1.0 +croniter==1.2.0 # via flytekit cryptography==36.0.1 - # via - # paramiko - # secretstorage + # via secretstorage dataclasses-json==0.5.6 # via flytekit -decorator==5.1.0 - # via - # ipython - # retry -defusedxml==0.7.1 - # via nbconvert +decorator==5.1.1 + # via retry deprecated==1.2.13 # via flytekit -diskcache==5.3.0 +diskcache==5.4.0 # via flytekit docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -entrypoints==0.3 - # via - # jupyter-client - # nbconvert - # papermill -flyteidl==0.21.17 +flyteidl==0.22.0 # via flytekit -gevent==21.12.0 - # via sagemaker-training -greenlet==1.1.2 - # via gevent grpcio==1.43.0 # via flytekit -hmsclient==0.1.1 - # via flytekit idna==3.3 # via requests -importlib-metadata==4.10.0 +importlib-metadata==4.10.1 # via keyring -inotify_simple==1.2.1 - # via sagemaker-training -ipykernel==5.5.6 - # via flytekit -ipython==7.31.1 - # via ipykernel -ipython-genutils==0.2.0 - # via - # ipykernel - # nbformat -jedi==0.18.1 - # via ipython jeepney==0.7.1 # via # keyring @@ -117,31 +66,11 @@ jinja2==3.0.3 # via # cookiecutter # jinja2-time - # nbconvert jinja2-time==0.2.0 # via cookiecutter -jmespath==0.10.0 - # via - # boto3 - # botocore jsonschema==3.2.0 - # via - # -r requirements.in - # nbformat -jupyter-client==7.1.0 - # via - # ipykernel - # nbclient -jupyter-core==4.9.1 - # via - # jupyter-client - # nbconvert - # nbformat -jupyterlab-pygments==0.1.2 - # via nbconvert -k8s-proto==0.0.3 - # via flytekit -keyring==23.4.0 + # via -r requirements.in +keyring==23.5.0 # via flytekit markupsafe==2.0.1 # via jinja2 @@ -154,100 +83,38 @@ marshmallow-enum==1.5.1 # via dataclasses-json marshmallow-jsonschema==0.13.0 # via flytekit -matplotlib-inline==0.1.3 - # via ipython -mistune==0.8.4 - # via nbconvert mypy-extensions==0.4.3 - # via - # black - # typing-inspect -natsort==8.0.2 - # via flytekit -nbclient==0.5.9 - # via - # nbconvert - # papermill -nbconvert==6.3.0 + # via typing-inspect +natsort==8.1.0 # via flytekit -nbformat==5.1.3 - # via - # nbclient - # nbconvert - # papermill -nest-asyncio==1.5.4 - # via - # jupyter-client - # nbclient numpy==1.21.5 # via - # flytekit + # -r requirements.in # pandas # pyarrow - # sagemaker-training - # scipy -packaging==21.3 - # via bleach pandas==1.3.5 - # via flytekit -pandocfilters==1.5.0 - # via nbconvert -papermill==2.3.3 - # via flytekit -paramiko==2.8.1 - # via sagemaker-training -parso==0.8.3 - # via jedi -pathspec==0.9.0 - # via black -pexpect==4.8.0 - # via ipython -pickleshare==0.7.5 - # via ipython -platformdirs==2.4.0 - # via black + # via + # -r requirements.in + # flytekit poyo==0.5.0 # via cookiecutter -prompt-toolkit==3.0.24 - # via ipython -protobuf==3.19.1 +protobuf==3.19.4 # via # flyteidl # flytekit - # k8s-proto - # sagemaker-training -psutil==5.8.0 - # via sagemaker-training -ptyprocess==0.7.0 - # via pexpect py==1.11.0 # via retry -py4j==0.10.9.2 - # via pyspark pyarrow==6.0.1 # via flytekit pycparser==2.21 # via cffi -pygments==2.10.0 - # via - # ipython - # jupyterlab-pygments - # nbconvert -pynacl==1.4.0 - # via paramiko -pyparsing==3.0.6 - # via packaging -pyrsistent==0.18.0 +pyrsistent==0.18.1 # via jsonschema -pyspark==3.2.0 - # via flytekit -python-dateutil==2.8.1 +python-dateutil==2.8.2 # via # arrow - # botocore # croniter # flytekit - # jupyter-client # pandas python-json-logger==2.0.2 # via flytekit @@ -260,110 +127,51 @@ pytz==2021.3 # flytekit # pandas pyyaml==5.4.1 - # via - # -r requirements.in - # papermill -pyzmq==22.3.0 - # via jupyter-client -regex==2021.11.10 + # via -r requirements.in +regex==2022.1.18 # via docker-image-py -requests==2.26.0 +requests==2.27.1 # via # cookiecutter # flytekit - # papermill # responses -responses==0.16.0 +responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit -retrying==1.3.3 - # via sagemaker-training -s3transfer==0.5.0 - # via boto3 -sagemaker-training==3.9.2 - # via flytekit -scipy==1.7.3 - # via sagemaker-training secretstorage==3.3.1 # via keyring six==1.16.0 # via - # bcrypt - # bleach # cookiecutter - # flytekit # grpcio # jsonschema - # pynacl # python-dateutil - # responses - # retrying - # sagemaker-training - # thrift sortedcontainers==2.4.0 # via flytekit statsd==3.3.0 # via flytekit -tenacity==8.0.1 - # via papermill -testpath==0.5.0 - # via nbconvert text-unidecode==1.3 # via python-slugify -textwrap3==0.9.2 - # via ansiwrap -thrift==0.15.0 - # via hmsclient -tomli==1.2.3 - # via black -tornado==6.1 - # via - # ipykernel - # jupyter-client -tqdm==4.62.3 - # via papermill -traitlets==5.1.1 - # via - # ipykernel - # ipython - # jupyter-client - # jupyter-core - # matplotlib-inline - # nbclient - # nbconvert - # nbformat typing-extensions==4.0.1 # via - # black + # flytekit # typing-inspect typing-inspect==0.7.1 # via dataclasses-json -urllib3==1.26.7 +urllib3==1.26.8 # via - # botocore # flytekit # requests # responses -wcwidth==0.2.5 - # via prompt-toolkit -webencodings==0.5.1 - # via bleach -werkzeug==2.0.2 - # via sagemaker-training wheel==0.37.1 # via flytekit wrapt==1.13.3 # via # deprecated # flytekit -zipp==3.6.0 +zipp==3.7.0 # via importlib-metadata -zope.event==4.5.0 - # via gevent -zope.interface==5.4.0 - # via gevent # The following packages are considered to be unsafe in a requirements file: -# pip # setuptools diff --git a/requirements.in b/requirements.in index 408185fca8..0cf01a6f13 100644 --- a/requirements.in +++ b/requirements.in @@ -7,3 +7,7 @@ attrs<21 # https://github.com/flyteorg/flyte/issues/1732. jsonschema<4 pyyaml<6 +# A number of dependencies, including pandas and numpy, started dropping support for python 3.7 in their latest +# releases (pandas>=1.4.0 and numpy>=1.22.0). More details in https://github.com/flyteorg/flyte/issues/2115. +pandas<1.4.0 +numpy<1.22.0 diff --git a/requirements.txt b/requirements.txt index 38f885cec8..d745fca2d9 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,107 +6,56 @@ # -e file:.#egg=flytekit # via -r requirements.in -ansiwrap==0.8.4 - # via papermill -arrow==1.2.1 +arrow==1.2.2 # via jinja2-time attrs==20.3.0 # via # -r requirements.in # jsonschema -backcall==0.2.0 - # via ipython -bcrypt==3.2.0 - # via paramiko binaryornot==0.4.4 # via cookiecutter -black==21.12b0 - # via papermill -bleach==4.1.0 - # via nbconvert -boto3==1.20.26 - # via sagemaker-training -botocore==1.23.26 - # via - # boto3 - # s3transfer certifi==2021.10.8 # via requests cffi==1.15.0 - # via - # bcrypt - # cryptography - # pynacl + # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.9 +charset-normalizer==2.0.11 # via requests checksumdir==1.2.0 # via flytekit -click==7.1.2 +click==8.0.3 # via - # black # cookiecutter # flytekit - # hmsclient - # papermill cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.1.0 +croniter==1.2.0 # via flytekit cryptography==36.0.1 - # via - # paramiko - # secretstorage + # via secretstorage dataclasses-json==0.5.6 # via flytekit -decorator==5.1.0 - # via - # ipython - # retry -defusedxml==0.7.1 - # via nbconvert +decorator==5.1.1 + # via retry deprecated==1.2.13 # via flytekit -diskcache==5.3.0 +diskcache==5.4.0 # via flytekit docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -entrypoints==0.3 - # via - # jupyter-client - # nbconvert - # papermill -flyteidl==0.21.17 +flyteidl==0.22.0 # via flytekit -gevent==21.12.0 - # via sagemaker-training -greenlet==1.1.2 - # via gevent grpcio==1.43.0 # via flytekit -hmsclient==0.1.1 - # via flytekit idna==3.3 # via requests -importlib-metadata==4.10.0 +importlib-metadata==4.10.1 # via keyring -inotify_simple==1.2.1 - # via sagemaker-training -ipykernel==5.5.6 - # via flytekit -ipython==7.31.1 - # via ipykernel -ipython-genutils==0.2.0 - # via - # ipykernel - # nbformat -jedi==0.18.1 - # via ipython jeepney==0.7.1 # via # keyring @@ -115,31 +64,11 @@ jinja2==3.0.3 # via # cookiecutter # jinja2-time - # nbconvert jinja2-time==0.2.0 # via cookiecutter -jmespath==0.10.0 - # via - # boto3 - # botocore jsonschema==3.2.0 - # via - # -r requirements.in - # nbformat -jupyter-client==7.1.0 - # via - # ipykernel - # nbclient -jupyter-core==4.9.1 - # via - # jupyter-client - # nbconvert - # nbformat -jupyterlab-pygments==0.1.2 - # via nbconvert -k8s-proto==0.0.3 - # via flytekit -keyring==23.4.0 + # via -r requirements.in +keyring==23.5.0 # via flytekit markupsafe==2.0.1 # via jinja2 @@ -152,100 +81,38 @@ marshmallow-enum==1.5.1 # via dataclasses-json marshmallow-jsonschema==0.13.0 # via flytekit -matplotlib-inline==0.1.3 - # via ipython -mistune==0.8.4 - # via nbconvert mypy-extensions==0.4.3 - # via - # black - # typing-inspect -natsort==8.0.2 - # via flytekit -nbclient==0.5.9 - # via - # nbconvert - # papermill -nbconvert==6.3.0 + # via typing-inspect +natsort==8.1.0 # via flytekit -nbformat==5.1.3 - # via - # nbclient - # nbconvert - # papermill -nest-asyncio==1.5.4 - # via - # jupyter-client - # nbclient numpy==1.21.5 # via - # flytekit + # -r requirements.in # pandas # pyarrow - # sagemaker-training - # scipy -packaging==21.3 - # via bleach pandas==1.3.5 - # via flytekit -pandocfilters==1.5.0 - # via nbconvert -papermill==2.3.3 - # via flytekit -paramiko==2.8.1 - # via sagemaker-training -parso==0.8.3 - # via jedi -pathspec==0.9.0 - # via black -pexpect==4.8.0 - # via ipython -pickleshare==0.7.5 - # via ipython -platformdirs==2.4.0 - # via black + # via + # -r requirements.in + # flytekit poyo==0.5.0 # via cookiecutter -prompt-toolkit==3.0.24 - # via ipython -protobuf==3.19.1 +protobuf==3.19.4 # via # flyteidl # flytekit - # k8s-proto - # sagemaker-training -psutil==5.8.0 - # via sagemaker-training -ptyprocess==0.7.0 - # via pexpect py==1.11.0 # via retry -py4j==0.10.9.2 - # via pyspark pyarrow==6.0.1 # via flytekit pycparser==2.21 # via cffi -pygments==2.10.0 - # via - # ipython - # jupyterlab-pygments - # nbconvert -pynacl==1.4.0 - # via paramiko -pyparsing==3.0.6 - # via packaging -pyrsistent==0.18.0 +pyrsistent==0.18.1 # via jsonschema -pyspark==3.2.0 - # via flytekit -python-dateutil==2.8.1 +python-dateutil==2.8.2 # via # arrow - # botocore # croniter # flytekit - # jupyter-client # pandas python-json-logger==2.0.2 # via flytekit @@ -258,110 +125,51 @@ pytz==2021.3 # flytekit # pandas pyyaml==5.4.1 - # via - # -r requirements.in - # papermill -pyzmq==22.3.0 - # via jupyter-client -regex==2021.11.10 + # via -r requirements.in +regex==2022.1.18 # via docker-image-py -requests==2.26.0 +requests==2.27.1 # via # cookiecutter # flytekit - # papermill # responses -responses==0.16.0 +responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit -retrying==1.3.3 - # via sagemaker-training -s3transfer==0.5.0 - # via boto3 -sagemaker-training==3.9.2 - # via flytekit -scipy==1.7.3 - # via sagemaker-training secretstorage==3.3.1 # via keyring six==1.16.0 # via - # bcrypt - # bleach # cookiecutter - # flytekit # grpcio # jsonschema - # pynacl # python-dateutil - # responses - # retrying - # sagemaker-training - # thrift sortedcontainers==2.4.0 # via flytekit statsd==3.3.0 # via flytekit -tenacity==8.0.1 - # via papermill -testpath==0.5.0 - # via nbconvert text-unidecode==1.3 # via python-slugify -textwrap3==0.9.2 - # via ansiwrap -thrift==0.15.0 - # via hmsclient -tomli==1.2.3 - # via black -tornado==6.1 - # via - # ipykernel - # jupyter-client -tqdm==4.62.3 - # via papermill -traitlets==5.1.1 - # via - # ipykernel - # ipython - # jupyter-client - # jupyter-core - # matplotlib-inline - # nbclient - # nbconvert - # nbformat typing-extensions==4.0.1 # via - # black + # flytekit # typing-inspect typing-inspect==0.7.1 # via dataclasses-json -urllib3==1.26.7 +urllib3==1.26.8 # via - # botocore # flytekit # requests # responses -wcwidth==0.2.5 - # via prompt-toolkit -webencodings==0.5.1 - # via bleach -werkzeug==2.0.2 - # via sagemaker-training wheel==0.37.1 # via flytekit wrapt==1.13.3 # via # deprecated # flytekit -zipp==3.6.0 +zipp==3.7.0 # via importlib-metadata -zope.event==4.5.0 - # via gevent -zope.interface==5.4.0 - # via gevent # The following packages are considered to be unsafe in a requirements file: -# pip # setuptools diff --git a/setup.py b/setup.py index 740a1d81e5..48250dac6b 100644 --- a/setup.py +++ b/setup.py @@ -39,7 +39,7 @@ "wheel>=0.30.0,<1.0.0", "pandas>=1.0.0,<2.0.0", "pyarrow>=4.0.0,<7.0.0", - "click>=6.6,<8.0", + "click>=6.6,<9.0", "croniter>=0.3.20,<4.0.0", "deprecated>=1.0,<2.0", "python-dateutil>=2.1", diff --git a/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt b/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt index 8fcc781de3..dc2ddebe81 100644 --- a/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt +++ b/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt @@ -4,7 +4,7 @@ # # make tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt # -arrow==1.2.1 +arrow==1.2.2 # via jinja2-time binaryornot==0.4.4 # via cookiecutter @@ -14,7 +14,7 @@ cffi==1.15.0 # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.9 +charset-normalizer==2.0.11 # via requests checksumdir==1.2.0 # via flytekit @@ -26,7 +26,7 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.1.0 +croniter==1.2.0 # via flytekit cryptography==36.0.1 # via secretstorage @@ -34,27 +34,27 @@ cycler==0.11.0 # via matplotlib dataclasses-json==0.5.6 # via flytekit -decorator==5.1.0 +decorator==5.1.1 # via retry deprecated==1.2.13 # via flytekit -diskcache==5.3.0 +diskcache==5.4.0 # via flytekit docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.21.17 +flyteidl==0.22.0 # via flytekit -flytekit==0.25.0 +flytekit==0.30.0 # via -r tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.in -fonttools==4.28.5 +fonttools==4.29.1 # via matplotlib grpcio==1.43.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.10.0 +importlib-metadata==4.10.1 # via keyring jeepney==0.7.1 # via @@ -68,7 +68,7 @@ jinja2-time==0.2.0 # via cookiecutter joblib==1.1.0 # via -r tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.in -keyring==23.4.0 +keyring==23.5.0 # via flytekit kiwisolver==1.3.2 # via matplotlib @@ -87,25 +87,25 @@ matplotlib==3.5.1 # via -r tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.in mypy-extensions==0.4.3 # via typing-inspect -natsort==8.0.2 +natsort==8.1.0 # via flytekit -numpy==1.21.5 +numpy==1.22.2 # via # matplotlib # opencv-python # pandas # pyarrow -opencv-python==4.5.4.60 +opencv-python==4.5.5.62 # via -r tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.in packaging==21.3 # via matplotlib -pandas==1.3.5 +pandas==1.4.0 # via flytekit -pillow==8.4.0 +pillow==9.0.1 # via matplotlib poyo==0.5.0 # via cookiecutter -protobuf==3.19.1 +protobuf==3.19.4 # via # flyteidl # flytekit @@ -115,11 +115,11 @@ pyarrow==6.0.1 # via flytekit pycparser==2.21 # via cffi -pyparsing==3.0.6 +pyparsing==3.0.7 # via # matplotlib # packaging -python-dateutil==2.8.1 +python-dateutil==2.8.2 # via # arrow # croniter @@ -132,18 +132,18 @@ python-slugify==5.0.2 # via cookiecutter pytimeparse==1.1.8 # via flytekit -pytz==2018.4 +pytz==2021.3 # via # flytekit # pandas -regex==2021.11.10 +regex==2022.1.18 # via docker-image-py -requests==2.26.0 +requests==2.27.1 # via # cookiecutter # flytekit # responses -responses==0.16.0 +responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit @@ -152,10 +152,8 @@ secretstorage==3.3.1 six==1.16.0 # via # cookiecutter - # flytekit # grpcio # python-dateutil - # responses sortedcontainers==2.4.0 # via flytekit statsd==3.3.0 @@ -163,10 +161,12 @@ statsd==3.3.0 text-unidecode==1.3 # via python-slugify typing-extensions==4.0.1 - # via typing-inspect + # via + # flytekit + # typing-inspect typing-inspect==0.7.1 # via dataclasses-json -urllib3==1.26.7 +urllib3==1.26.8 # via # flytekit # requests @@ -179,5 +179,5 @@ wrapt==1.13.3 # via # deprecated # flytekit -zipp==3.6.0 +zipp==3.7.0 # via importlib-metadata From 30ef9e95a9636affc8f580441b3a0ec70f1c5f5f Mon Sep 17 00:00:00 2001 From: Kenny Workman <31255434+kennyworkman@users.noreply.github.com> Date: Tue, 8 Feb 2022 11:36:47 -0800 Subject: [PATCH 078/128] TypeAnnotation (#759) * feat: support for annotated simple + list Signed-off-by: Kenny Workman * feat: addition of annotation att to Signed-off-by: Kenny Workman * feat: core obj Signed-off-by: Kenny Workman * feat: proto model Signed-off-by: Kenny Workman * feat: testing suite Signed-off-by: Kenny Workman * fix: more stable typing introspection Signed-off-by: Kenny Workman * fix: strip legacy Signed-off-by: Kenny Workman * fix: explicitly allow only one annotation Signed-off-by: Kenny Workman * feat: direct type transformer tests Signed-off-by: Kenny Workman * fix: there and back test Signed-off-by: Kenny Workman * fix: typing_extensions for get_origin Signed-off-by: Kenny Workman * fix: more semantic list generic unwrap Signed-off-by: Kenny Workman * fix: tmp requirements file with custom idl Signed-off-by: Kenny Workman * fix: nits Signed-off-by: Kenny Workman * feat: semantic error for unsupported complex literals Signed-off-by: Kenny Workman * fix: but Signed-off-by: Kenny Workman * feat: more tests ;) Signed-off-by: Kenny Workman * fix: imports Signed-off-by: Kenny Workman * fix: complex annotations Signed-off-by: Kenny Workman * fix: temp requirements files for unit tests Signed-off-by: Kenny Workman * fix: lint bug Signed-off-by: Kenny Workman * fix: tmp setup.py Signed-off-by: Kenny Workman * fix: use typing_extensions Signed-off-by: Kenny Workman * fix: typing_extensions for annotated Signed-off-by: Kenny Workman * fix: typing_ext Signed-off-by: Kenny Workman * fix: plugin tmp requirements Signed-off-by: Kenny Workman * fix: bump requirements Signed-off-by: Kenny Workman * fix: doc requirements Signed-off-by: Kenny Workman * fix: whitespace Signed-off-by: Kenny Workman * fix: bump flytekit Signed-off-by: Kenny Workman * fix: numpy version Signed-off-by: Kenny Workman * fix: lint Signed-off-by: Kenny Workman * fix: pandas version Signed-off-by: Kenny Workman * fix: bump requirements Signed-off-by: Kenny Workman * fix: test import Signed-off-by: Kenny Workman * fix: flake8 lint Signed-off-by: Kenny Workman * fix: merge Signed-off-by: Kenny Workman * fix: requirements Signed-off-by: Kenny Workman * fix: requirements Signed-off-by: Kenny Workman * fix: lint Signed-off-by: Kenny Workman * fix: papermill req Signed-off-by: Kenny Workman * fix: req Signed-off-by: Kenny Workman --- .gitignore | 2 +- dev-requirements.txt | 1 - flytekit/core/annotation.py | 30 ++++++++ flytekit/core/type_engine.py | 74 +++++++++++++++++-- flytekit/models/annotation.py | 43 +++++++++++ flytekit/models/types.py | 20 +++++ plugins/flytekit-aws-athena/requirements.txt | 8 -- .../flytekit-aws-sagemaker/requirements.txt | 6 -- plugins/flytekit-data-fsspec/requirements.txt | 8 -- plugins/flytekit-dolt/requirements.txt | 8 -- .../requirements.txt | 6 -- plugins/flytekit-hive/requirements.txt | 8 -- plugins/flytekit-k8s-pod/requirements.txt | 8 -- plugins/flytekit-kf-mpi/requirements.txt | 8 -- plugins/flytekit-kf-pytorch/requirements.txt | 8 -- .../flytekit-kf-tensorflow/requirements.txt | 8 -- plugins/flytekit-modin/requirements.txt | 8 -- plugins/flytekit-pandera/requirements.txt | 8 -- .../flytekit-papermill/dev-requirements.txt | 2 +- plugins/flytekit-papermill/requirements.txt | 2 +- plugins/flytekit-snowflake/requirements.txt | 8 -- plugins/flytekit-spark/requirements.txt | 8 -- plugins/flytekit-sqlalchemy/requirements.txt | 8 -- setup.py | 3 +- .../workflows/requirements.txt | 11 --- tests/flytekit/unit/core/test_type_engine.py | 62 ++++++++++++++++ .../unit/core/test_typing_annotation.py | 62 ++++++++++++++++ tests/flytekit/unit/models/test_types.py | 11 +++ 28 files changed, 300 insertions(+), 139 deletions(-) create mode 100644 flytekit/core/annotation.py create mode 100644 flytekit/models/annotation.py create mode 100644 tests/flytekit/unit/core/test_typing_annotation.py diff --git a/.gitignore b/.gitignore index ec56a0f239..73c81d2a6a 100644 --- a/.gitignore +++ b/.gitignore @@ -3,7 +3,7 @@ *.pyt *.pytc *.egg-info -.*.swp +.*.sw* .DS_Store venv/ .venv/ diff --git a/dev-requirements.txt b/dev-requirements.txt index 28c22afa15..6e17840325 100644 --- a/dev-requirements.txt +++ b/dev-requirements.txt @@ -32,7 +32,6 @@ certifi==2021.10.8 # requests cffi==1.15.0 # via - # -c requirements.txt # bcrypt # cryptography # pynacl diff --git a/flytekit/core/annotation.py b/flytekit/core/annotation.py new file mode 100644 index 0000000000..b4a70a6469 --- /dev/null +++ b/flytekit/core/annotation.py @@ -0,0 +1,30 @@ +from typing import Any, Dict + + +class FlyteAnnotation: + """A core object to add arbitrary annotations to flyte types. + + This metadata is ingested as a python dictionary and will be serialized + into fields on the flyteidl type literals. This data is not accessible at + runtime but rather can be retrieved from flyteadmin for custom presentation + of typed parameters. + + Flytekit expects to receive a maximum of one `FlyteAnnotation` object + within each typehint. + + For a task definition: + + .. code-block:: python + + @task + def x(a: typing.Annotated[int, FlyteAnnotation({"foo": {"bar": 1}})]): + return + + """ + + def __init__(self, data: Dict[str, Any]): + self._data = data + + @property + def data(self): + return self._data diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 2d41512c8e..c8e2ce354e 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -26,12 +26,14 @@ from marshmallow_enum import EnumField, LoadDumpOptions from marshmallow_jsonschema import JSONSchema +from flytekit.core.annotation import FlyteAnnotation from flytekit.core.context_manager import FlyteContext from flytekit.core.type_helpers import load_type_from_tag from flytekit.exceptions import user as user_exceptions from flytekit.loggers import logger from flytekit.models import interface as _interface_models from flytekit.models import types as _type_models +from flytekit.models.annotation import TypeAnnotation as TypeAnnotationModel from flytekit.models.core import types as _core_types from flytekit.models.literals import ( Blob, @@ -268,6 +270,12 @@ def get_literal_type(self, t: Type[T]) -> LiteralType: Extracts the Literal type definition for a Dataclass and returns a type Struct. If possible also extracts the JSONSchema for the dataclass. """ + if get_origin(t) is Annotated: + raise ValueError( + "Flytekit does not currently have support for FlyteAnnotations applied to Dataclass." + f"Type {t} cannot be parsed." + ) + if not issubclass(t, DataClassJsonMixin): raise AssertionError( f"Dataclass {t} should be decorated with @dataclass_json to be " f"serialized correctly" @@ -587,6 +595,7 @@ def get_transformer(cls, python_type: Type) -> TypeTransformer[T]: TODO lets make this deterministic by using an ordered dict """ + # Step 1 if get_origin(python_type) is Annotated: python_type = get_args(python_type)[0] @@ -596,8 +605,14 @@ def get_transformer(cls, python_type: Type) -> TypeTransformer[T]: # Step 2 if hasattr(python_type, "__origin__"): + # Handling of annotated generics, eg: + # Annotated[typing.List[int], 'foo'] + if get_origin(python_type) is Annotated: + return cls.get_transformer(get_args(python_type)[0]) + if python_type.__origin__ in cls._REGISTRY: return cls._REGISTRY[python_type.__origin__] + raise ValueError(f"Generic Type {python_type.__origin__} not supported currently in Flytekit.") # Step 3 @@ -626,7 +641,30 @@ def to_literal_type(cls, python_type: Type) -> LiteralType: Converts a python type into a flyte specific ``LiteralType`` """ transformer = cls.get_transformer(python_type) - return transformer.get_literal_type(python_type) + res = transformer.get_literal_type(python_type) + data = None + if get_origin(python_type) is Annotated: + for x in get_args(python_type)[1:]: + if not isinstance(x, FlyteAnnotation): + continue + if data is not None: + raise ValueError( + f"More than one FlyteAnnotation used within {python_type} typehint. Flytekit requires a max of one." + ) + data = x.data + if data is not None: + idl_type_annotation = TypeAnnotationModel(annotations=data) + return LiteralType( + simple=res.simple, + schema=res.schema, + collection_type=res.collection_type, + map_value_type=res.map_value_type, + blob=res.blob, + enum_type=res.enum_type, + metadata=res.metadata, + annotation=idl_type_annotation, + ) + return res @classmethod def to_literal(cls, ctx: FlyteContext, python_val: typing.Any, python_type: Type, expected: LiteralType) -> Literal: @@ -754,9 +792,16 @@ def get_sub_type(t: Type[T]) -> Type[T]: """ Return the generic Type T of the List """ - if hasattr(t, "__origin__") and t.__origin__ is list: # type: ignore - if hasattr(t, "__args__"): - return t.__args__[0] # type: ignore + + if hasattr(t, "__origin__"): + # Handle annotation on list generic, eg: + # Annotated[typing.List[int], 'foo'] + if get_origin(t) is Annotated: + return ListTransformer.get_sub_type(get_args(t)[0]) + + if t.__origin__ is list and hasattr(t, "__args__"): + return t.__args__[0] + raise ValueError("Only generic univariate typing.List[T] type is supported.") def get_literal_type(self, t: Type[T]) -> Optional[LiteralType]: @@ -898,9 +943,17 @@ def get_dict_types(t: Optional[Type[dict]]) -> typing.Tuple[Optional[type], Opti """ Return the generic Type T of the Dict """ - if hasattr(t, "__origin__") and t.__origin__ is dict: # type: ignore - if hasattr(t, "__args__"): - return t.__args__ # type: ignore + _origin = get_origin(t) + _args = get_args(t) + if _origin is not None: + if _origin is Annotated: + raise ValueError( + f"Flytekit does not currently have support \ + for FlyteAnnotations applied to dicts. {t} cannot be \ + parsed." + ) + if _origin is dict and _args is not None: + return _args return None, None @staticmethod @@ -1051,6 +1104,13 @@ def __init__(self): super().__init__(name="DefaultEnumTransformer", t=enum.Enum) def get_literal_type(self, t: Type[T]) -> LiteralType: + if get_origin(t) is Annotated: + raise ValueError( + f"Flytekit does not currently have support \ + for FlyteAnnotations applied to enums. {t} cannot be \ + parsed." + ) + values = [v.value for v in t] # type: ignore if not isinstance(values[0], str): raise TypeTransformerFailedError("Only EnumTypes with value of string are supported") diff --git a/flytekit/models/annotation.py b/flytekit/models/annotation.py new file mode 100644 index 0000000000..ced935c57e --- /dev/null +++ b/flytekit/models/annotation.py @@ -0,0 +1,43 @@ +import json as _json +from typing import Any, Dict + +from flyteidl.core import types_pb2 as _types_pb2 +from google.protobuf import json_format as _json_format +from google.protobuf import struct_pb2 as _struct + + +class TypeAnnotation: + """Python class representation of the flyteidl TypeAnnotation message.""" + + def __init__(self, annotations: Dict[str, Any]): + self._annotations = annotations + + @property + def annotations(self) -> Dict[str, Any]: + """ + :rtype: dict[str, Any] + """ + return self._annotations + + def to_flyte_idl(self) -> _types_pb2.TypeAnnotation: + """ + :rtype: flyteidl.core.types_pb2.TypeAnnotation + """ + + if self._annotations is not None: + annotations = _json_format.Parse(_json.dumps(self.annotations), _struct.Struct()) + else: + annotations = None + + return _types_pb2.TypeAnnotation( + annotations=annotations, + ) + + @classmethod + def from_flyte_idl(cls, proto): + """ + :param flyteidl.core.types_pb2.TypeAnnotation proto: + :rtype: TypeAnnotation + """ + + return cls(annotations=_json_format.MessageToDict(proto.annotations)) diff --git a/flytekit/models/types.py b/flytekit/models/types.py index 50c3432838..c8cc1307d3 100644 --- a/flytekit/models/types.py +++ b/flytekit/models/types.py @@ -6,6 +6,7 @@ from google.protobuf import struct_pb2 as _struct from flytekit.models import common as _common +from flytekit.models.annotation import TypeAnnotation as TypeAnnotationModel from flytekit.models.core import types as _core_types @@ -241,6 +242,7 @@ def __init__( structure=None, structured_dataset_type=None, metadata=None, + annotation=None, ): """ This is a oneof message, only one of the kwargs may be set, representing one of the Flyte types. @@ -256,6 +258,8 @@ def __init__( :param flytekit.models.core.types.TypeStructure structure: Type matching hints :param flytekit.models.core.types.StructuredDatasetType structured_dataset_type: structured dataset :param dict[Text, T] metadata: Additional data describing the type + :param flytekit.models.annotation.FlyteAnnotation annotation: Additional data + describing the type _intended to be saturated by the client_ """ self._simple = simple self._schema = schema @@ -267,6 +271,7 @@ def __init__( self._structure = structure self._structured_dataset_type = structured_dataset_type self._metadata = metadata + self._annotation = annotation @property def simple(self) -> SimpleType: @@ -316,18 +321,31 @@ def metadata(self): """ return self._metadata + @property + def annotation(self) -> TypeAnnotationModel: + """ + :rtype: flytekit.models.annotation.TypeAnnotation + """ + return self._annotation + @metadata.setter def metadata(self, value): self._metadata = value + @annotation.setter + def annotation(self, value): + self.annotation = value + def to_flyte_idl(self): """ :rtype: flyteidl.core.types_pb2.LiteralType """ + if self.metadata is not None: metadata = _json_format.Parse(_json.dumps(self.metadata), _struct.Struct()) else: metadata = None + t = _types_pb2.LiteralType( simple=self.simple if self.simple is not None else None, schema=self.schema.to_flyte_idl() if self.schema is not None else None, @@ -341,6 +359,7 @@ def to_flyte_idl(self): if self.structured_dataset_type else None, metadata=metadata, + annotation=self.annotation.to_flyte_idl() if self.annotation else None, ) return t @@ -369,6 +388,7 @@ def from_flyte_idl(cls, proto): if proto.HasField("structured_dataset_type") else None, metadata=_json_format.MessageToDict(proto.metadata) or None, + annotation=TypeAnnotationModel.from_flyte_idl(proto.annotation) if proto.HasField("annotation") else None, ) diff --git a/plugins/flytekit-aws-athena/requirements.txt b/plugins/flytekit-aws-athena/requirements.txt index 6fe91c53c8..a1acae494f 100644 --- a/plugins/flytekit-aws-athena/requirements.txt +++ b/plugins/flytekit-aws-athena/requirements.txt @@ -12,8 +12,6 @@ binaryornot==0.4.4 # via cookiecutter certifi==2021.10.8 # via requests -cffi==1.15.0 - # via cryptography chardet==4.0.0 # via binaryornot charset-normalizer==2.0.11 @@ -54,10 +52,6 @@ idna==3.3 # via requests importlib-metadata==4.10.1 # via keyring -jeepney==0.7.1 - # via - # keyring - # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -126,8 +120,6 @@ responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit -secretstorage==3.3.1 - # via keyring six==1.16.0 # via # cookiecutter diff --git a/plugins/flytekit-aws-sagemaker/requirements.txt b/plugins/flytekit-aws-sagemaker/requirements.txt index 2624f28a98..11396f213e 100644 --- a/plugins/flytekit-aws-sagemaker/requirements.txt +++ b/plugins/flytekit-aws-sagemaker/requirements.txt @@ -73,10 +73,6 @@ importlib-metadata==4.10.1 # via keyring inotify-simple==1.2.1 # via sagemaker-training -jeepney==0.7.1 - # via - # keyring - # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -167,8 +163,6 @@ sagemaker-training==3.9.2 # via flytekitplugins-awssagemaker scipy==1.8.0 # via sagemaker-training -secretstorage==3.3.1 - # via keyring six==1.16.0 # via # bcrypt diff --git a/plugins/flytekit-data-fsspec/requirements.txt b/plugins/flytekit-data-fsspec/requirements.txt index a14a5bf6cf..c644ac98e7 100644 --- a/plugins/flytekit-data-fsspec/requirements.txt +++ b/plugins/flytekit-data-fsspec/requirements.txt @@ -12,8 +12,6 @@ binaryornot==0.4.4 # via cookiecutter certifi==2021.10.8 # via requests -cffi==1.15.0 - # via cryptography chardet==4.0.0 # via binaryornot charset-normalizer==2.0.11 @@ -56,10 +54,6 @@ idna==3.3 # via requests importlib-metadata==4.10.1 # via keyring -jeepney==0.7.1 - # via - # keyring - # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -128,8 +122,6 @@ responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit -secretstorage==3.3.1 - # via keyring six==1.16.0 # via # cookiecutter diff --git a/plugins/flytekit-dolt/requirements.txt b/plugins/flytekit-dolt/requirements.txt index 559d5b2085..01dfe2cb6b 100644 --- a/plugins/flytekit-dolt/requirements.txt +++ b/plugins/flytekit-dolt/requirements.txt @@ -12,8 +12,6 @@ binaryornot==0.4.4 # via cookiecutter certifi==2021.10.8 # via requests -cffi==1.15.0 - # via cryptography chardet==4.0.0 # via binaryornot charset-normalizer==2.0.11 @@ -60,10 +58,6 @@ idna==3.3 # via requests importlib-metadata==4.10.1 # via keyring -jeepney==0.7.1 - # via - # keyring - # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -134,8 +128,6 @@ responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit -secretstorage==3.3.1 - # via keyring six==1.16.0 # via # cookiecutter diff --git a/plugins/flytekit-greatexpectations/requirements.txt b/plugins/flytekit-greatexpectations/requirements.txt index fdbb857264..624f3dc633 100644 --- a/plugins/flytekit-greatexpectations/requirements.txt +++ b/plugins/flytekit-greatexpectations/requirements.txt @@ -110,10 +110,6 @@ ipywidgets==7.6.5 # via great-expectations jedi==0.18.1 # via ipython -jeepney==0.7.1 - # via - # keyring - # secretstorage jinja2==3.0.3 # via # altair @@ -295,8 +291,6 @@ ruamel.yaml.clib==0.2.6 # via ruamel.yaml scipy==1.7.3 # via great-expectations -secretstorage==3.3.1 - # via keyring send2trash==1.8.0 # via notebook six==1.16.0 diff --git a/plugins/flytekit-hive/requirements.txt b/plugins/flytekit-hive/requirements.txt index 269d52fceb..3ae057f9a0 100644 --- a/plugins/flytekit-hive/requirements.txt +++ b/plugins/flytekit-hive/requirements.txt @@ -12,8 +12,6 @@ binaryornot==0.4.4 # via cookiecutter certifi==2021.10.8 # via requests -cffi==1.15.0 - # via cryptography chardet==4.0.0 # via binaryornot charset-normalizer==2.0.11 @@ -54,10 +52,6 @@ idna==3.3 # via requests importlib-metadata==4.10.1 # via keyring -jeepney==0.7.1 - # via - # keyring - # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -126,8 +120,6 @@ responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit -secretstorage==3.3.1 - # via keyring six==1.16.0 # via # cookiecutter diff --git a/plugins/flytekit-k8s-pod/requirements.txt b/plugins/flytekit-k8s-pod/requirements.txt index 6574f4ab66..d60d9c46b6 100644 --- a/plugins/flytekit-k8s-pod/requirements.txt +++ b/plugins/flytekit-k8s-pod/requirements.txt @@ -16,8 +16,6 @@ certifi==2021.10.8 # via # kubernetes # requests -cffi==1.15.0 - # via cryptography chardet==4.0.0 # via binaryornot charset-normalizer==2.0.11 @@ -60,10 +58,6 @@ idna==3.3 # via requests importlib-metadata==4.10.1 # via keyring -jeepney==0.7.1 - # via - # keyring - # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -151,8 +145,6 @@ retry==0.9.2 # via flytekit rsa==4.8 # via google-auth -secretstorage==3.3.1 - # via keyring six==1.16.0 # via # cookiecutter diff --git a/plugins/flytekit-kf-mpi/requirements.txt b/plugins/flytekit-kf-mpi/requirements.txt index ef91243daf..c13f9f3b74 100644 --- a/plugins/flytekit-kf-mpi/requirements.txt +++ b/plugins/flytekit-kf-mpi/requirements.txt @@ -12,8 +12,6 @@ binaryornot==0.4.4 # via cookiecutter certifi==2021.10.8 # via requests -cffi==1.15.0 - # via cryptography chardet==4.0.0 # via binaryornot charset-normalizer==2.0.11 @@ -56,10 +54,6 @@ idna==3.3 # via requests importlib-metadata==4.10.1 # via keyring -jeepney==0.7.1 - # via - # keyring - # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -128,8 +122,6 @@ responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit -secretstorage==3.3.1 - # via keyring six==1.16.0 # via # cookiecutter diff --git a/plugins/flytekit-kf-pytorch/requirements.txt b/plugins/flytekit-kf-pytorch/requirements.txt index 9fc8665f43..1696646c87 100644 --- a/plugins/flytekit-kf-pytorch/requirements.txt +++ b/plugins/flytekit-kf-pytorch/requirements.txt @@ -12,8 +12,6 @@ binaryornot==0.4.4 # via cookiecutter certifi==2021.10.8 # via requests -cffi==1.15.0 - # via cryptography chardet==4.0.0 # via binaryornot charset-normalizer==2.0.11 @@ -54,10 +52,6 @@ idna==3.3 # via requests importlib-metadata==4.10.1 # via keyring -jeepney==0.7.1 - # via - # keyring - # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -126,8 +120,6 @@ responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit -secretstorage==3.3.1 - # via keyring six==1.16.0 # via # cookiecutter diff --git a/plugins/flytekit-kf-tensorflow/requirements.txt b/plugins/flytekit-kf-tensorflow/requirements.txt index ae428e9215..d3b6253664 100644 --- a/plugins/flytekit-kf-tensorflow/requirements.txt +++ b/plugins/flytekit-kf-tensorflow/requirements.txt @@ -12,8 +12,6 @@ binaryornot==0.4.4 # via cookiecutter certifi==2021.10.8 # via requests -cffi==1.15.0 - # via cryptography chardet==4.0.0 # via binaryornot charset-normalizer==2.0.11 @@ -54,10 +52,6 @@ idna==3.3 # via requests importlib-metadata==4.10.1 # via keyring -jeepney==0.7.1 - # via - # keyring - # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -126,8 +120,6 @@ responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit -secretstorage==3.3.1 - # via keyring six==1.16.0 # via # cookiecutter diff --git a/plugins/flytekit-modin/requirements.txt b/plugins/flytekit-modin/requirements.txt index adaa0b7307..4f3b01e5e4 100644 --- a/plugins/flytekit-modin/requirements.txt +++ b/plugins/flytekit-modin/requirements.txt @@ -16,8 +16,6 @@ binaryornot==0.4.4 # via cookiecutter certifi==2021.10.8 # via requests -cffi==1.15.0 - # via cryptography chardet==4.0.0 # via binaryornot charset-normalizer==2.0.11 @@ -69,10 +67,6 @@ idna==3.3 # via requests importlib-metadata==4.10.1 # via keyring -jeepney==0.7.1 - # via - # keyring - # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -166,8 +160,6 @@ responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit -secretstorage==3.3.1 - # via keyring six==1.16.0 # via # cookiecutter diff --git a/plugins/flytekit-pandera/requirements.txt b/plugins/flytekit-pandera/requirements.txt index 82df51d30e..f0b65cec17 100644 --- a/plugins/flytekit-pandera/requirements.txt +++ b/plugins/flytekit-pandera/requirements.txt @@ -12,8 +12,6 @@ binaryornot==0.4.4 # via cookiecutter certifi==2021.10.8 # via requests -cffi==1.15.0 - # via cryptography chardet==4.0.0 # via binaryornot charset-normalizer==2.0.11 @@ -54,10 +52,6 @@ idna==3.3 # via requests importlib-metadata==4.10.1 # via keyring -jeepney==0.7.1 - # via - # keyring - # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -137,8 +131,6 @@ responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit -secretstorage==3.3.1 - # via keyring six==1.16.0 # via # cookiecutter diff --git a/plugins/flytekit-papermill/dev-requirements.txt b/plugins/flytekit-papermill/dev-requirements.txt index 849cef4c9f..d647708f20 100644 --- a/plugins/flytekit-papermill/dev-requirements.txt +++ b/plugins/flytekit-papermill/dev-requirements.txt @@ -38,7 +38,7 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.21.23 +flyteidl==0.22.0 # via flytekit flytekit==0.26.0 # via flytekitplugins-spark diff --git a/plugins/flytekit-papermill/requirements.txt b/plugins/flytekit-papermill/requirements.txt index b06adde647..8c3e4570e9 100644 --- a/plugins/flytekit-papermill/requirements.txt +++ b/plugins/flytekit-papermill/requirements.txt @@ -71,7 +71,7 @@ entrypoints==0.3 # papermill executing==0.8.2 # via stack-data -flyteidl==0.21.24 +flyteidl==0.22.0 # via flytekit flytekit==0.26.1 # via flytekitplugins-papermill diff --git a/plugins/flytekit-snowflake/requirements.txt b/plugins/flytekit-snowflake/requirements.txt index b563e54814..53951100b5 100644 --- a/plugins/flytekit-snowflake/requirements.txt +++ b/plugins/flytekit-snowflake/requirements.txt @@ -12,8 +12,6 @@ binaryornot==0.4.4 # via cookiecutter certifi==2021.10.8 # via requests -cffi==1.15.0 - # via cryptography chardet==4.0.0 # via binaryornot charset-normalizer==2.0.11 @@ -54,10 +52,6 @@ idna==3.3 # via requests importlib-metadata==4.10.1 # via keyring -jeepney==0.7.1 - # via - # keyring - # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -126,8 +120,6 @@ responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit -secretstorage==3.3.1 - # via keyring six==1.16.0 # via # cookiecutter diff --git a/plugins/flytekit-spark/requirements.txt b/plugins/flytekit-spark/requirements.txt index fed298d4b4..75d1e32d53 100644 --- a/plugins/flytekit-spark/requirements.txt +++ b/plugins/flytekit-spark/requirements.txt @@ -12,8 +12,6 @@ binaryornot==0.4.4 # via cookiecutter certifi==2021.10.8 # via requests -cffi==1.15.0 - # via cryptography chardet==4.0.0 # via binaryornot charset-normalizer==2.0.11 @@ -54,10 +52,6 @@ idna==3.3 # via requests importlib-metadata==4.10.1 # via keyring -jeepney==0.7.1 - # via - # keyring - # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -130,8 +124,6 @@ responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit -secretstorage==3.3.1 - # via keyring six==1.16.0 # via # cookiecutter diff --git a/plugins/flytekit-sqlalchemy/requirements.txt b/plugins/flytekit-sqlalchemy/requirements.txt index 134f8960e8..93dcc2790c 100644 --- a/plugins/flytekit-sqlalchemy/requirements.txt +++ b/plugins/flytekit-sqlalchemy/requirements.txt @@ -12,8 +12,6 @@ binaryornot==0.4.4 # via cookiecutter certifi==2021.10.8 # via requests -cffi==1.15.0 - # via cryptography chardet==4.0.0 # via binaryornot charset-normalizer==2.0.11 @@ -56,10 +54,6 @@ idna==3.3 # via requests importlib-metadata==4.10.1 # via keyring -jeepney==0.7.1 - # via - # keyring - # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -128,8 +122,6 @@ responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit -secretstorage==3.3.1 - # via keyring six==1.16.0 # via # cookiecutter diff --git a/setup.py b/setup.py index 48250dac6b..885f32b897 100644 --- a/setup.py +++ b/setup.py @@ -35,7 +35,7 @@ ] }, install_requires=[ - "flyteidl>=0.21.17", + "flyteidl>=0.22.0", "wheel>=0.30.0,<1.0.0", "pandas>=1.0.0,<2.0.0", "pyarrow>=4.0.0,<7.0.0", @@ -67,6 +67,7 @@ "checksumdir>=1.2.0", "cloudpickle>=2.0.0", "cookiecutter>=1.7.3", + "numpy<=1.22.1; python_version < '3.8.0'", ], extras_require=extras_require, scripts=[ diff --git a/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt b/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt index dc2ddebe81..93514898b2 100644 --- a/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt +++ b/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt @@ -10,8 +10,6 @@ binaryornot==0.4.4 # via cookiecutter certifi==2021.10.8 # via requests -cffi==1.15.0 - # via cryptography chardet==4.0.0 # via binaryornot charset-normalizer==2.0.11 @@ -28,8 +26,6 @@ cookiecutter==1.7.3 # via flytekit croniter==1.2.0 # via flytekit -cryptography==36.0.1 - # via secretstorage cycler==0.11.0 # via matplotlib dataclasses-json==0.5.6 @@ -56,10 +52,6 @@ idna==3.3 # via requests importlib-metadata==4.10.1 # via keyring -jeepney==0.7.1 - # via - # keyring - # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -121,7 +113,6 @@ pyparsing==3.0.7 # packaging python-dateutil==2.8.2 # via - # arrow # croniter # flytekit # matplotlib @@ -147,8 +138,6 @@ responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit -secretstorage==3.3.1 - # via keyring six==1.16.0 # via # cookiecutter diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index a046b58ee9..7de44b9935 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -9,6 +9,7 @@ import pandas as pd import pyarrow as pa import pytest +import typing_extensions from dataclasses_json import DataClassJsonMixin, dataclass_json from flyteidl.core import errors_pb2 from google.protobuf import json_format as _json_format @@ -21,6 +22,7 @@ from flytekit import kwtypes from flytekit.common.exceptions import user as user_exceptions from flytekit.common.types import primitives +from flytekit.core.annotation import FlyteAnnotation from flytekit.core.context_manager import FlyteContext, FlyteContextManager from flytekit.core.type_engine import ( DataclassTransformer, @@ -36,6 +38,7 @@ ) from flytekit.exceptions import user as user_exceptions from flytekit.models import types as model_types +from flytekit.models.annotation import TypeAnnotation from flytekit.models.core.types import BlobType from flytekit.models.literals import Blob, BlobMetadata, Literal, LiteralCollection, LiteralMap, Primitive, Scalar, Void from flytekit.models.types import LiteralType, SimpleType, TypeStructure @@ -1067,6 +1070,65 @@ def test_dict_to_literal_map_with_wrong_input_type(): TypeEngine.dict_to_literal_map(ctx, input, guessed_python_types) +def test_annotated_simple_types(): + def _check_annotation(t, annotation): + lt = TypeEngine.to_literal_type(t) + assert isinstance(lt.annotation, TypeAnnotation) + assert lt.annotation.annotations == annotation + + _check_annotation(typing_extensions.Annotated[int, FlyteAnnotation({"foo": "bar"})], {"foo": "bar"}) + _check_annotation(typing_extensions.Annotated[int, FlyteAnnotation(["foo", "bar"])], ["foo", "bar"]) + _check_annotation( + typing_extensions.Annotated[int, FlyteAnnotation({"d": {"test": "data"}, "l": ["nested", ["list"]]})], + {"d": {"test": "data"}, "l": ["nested", ["list"]]}, + ) + _check_annotation( + typing_extensions.Annotated[int, FlyteAnnotation(InnerStruct(a=1, b="fizz", c=[1]))], + InnerStruct(a=1, b="fizz", c=[1]), + ) + + +def test_annotated_list(): + t = typing_extensions.Annotated[typing.List[int], FlyteAnnotation({"foo": "bar"})] + lt = TypeEngine.to_literal_type(t) + assert isinstance(lt.annotation, TypeAnnotation) + assert lt.annotation.annotations == {"foo": "bar"} + + t = typing.List[typing_extensions.Annotated[int, FlyteAnnotation({"foo": "bar"})]] + lt = TypeEngine.to_literal_type(t) + assert isinstance(lt.collection_type.annotation, TypeAnnotation) + assert lt.collection_type.annotation.annotations == {"foo": "bar"} + + +def test_type_alias(): + inner_t = typing_extensions.Annotated[int, FlyteAnnotation("foo")] + t = typing_extensions.Annotated[inner_t, FlyteAnnotation("bar")] + with pytest.raises(ValueError): + TypeEngine.to_literal_type(t) + + +def test_unsupported_complex_literals(): + t = typing_extensions.Annotated[typing.Dict[int, str], FlyteAnnotation({"foo": "bar"})] + with pytest.raises(ValueError): + TypeEngine.to_literal_type(t) + + # Enum. + t = typing_extensions.Annotated[Color, FlyteAnnotation({"foo": "bar"})] + with pytest.raises(ValueError): + TypeEngine.to_literal_type(t) + + # Dataclass. + t = typing_extensions.Annotated[Result, FlyteAnnotation({"foo": "bar"})] + with pytest.raises(ValueError): + TypeEngine.to_literal_type(t) + + +def test_multiple_annotations(): + t = typing_extensions.Annotated[int, FlyteAnnotation({"foo": "bar"}), FlyteAnnotation({"anotha": "one"})] + with pytest.raises(Exception): + TypeEngine.to_literal_type(t) + + TestSchema = FlyteSchema[kwtypes(some_str=str)] diff --git a/tests/flytekit/unit/core/test_typing_annotation.py b/tests/flytekit/unit/core/test_typing_annotation.py new file mode 100644 index 0000000000..f999c62612 --- /dev/null +++ b/tests/flytekit/unit/core/test_typing_annotation.py @@ -0,0 +1,62 @@ +import typing +from collections import OrderedDict + +import typing_extensions + +from flytekit.core import context_manager +from flytekit.core.annotation import FlyteAnnotation +from flytekit.core.context_manager import Image, ImageConfig +from flytekit.core.task import task +from flytekit.models.annotation import TypeAnnotation +from flytekit.tools.translator import get_serializable + +default_img = Image(name="default", fqn="test", tag="tag") +serialization_settings = context_manager.SerializationSettings( + project="project", + domain="domain", + version="version", + env=None, + image_config=ImageConfig(default_image=default_img, images=[default_img]), +) +entity_mapping = OrderedDict() + + +@task +def x(a: typing_extensions.Annotated[int, FlyteAnnotation({"foo": {"bar": 1}})], b: str): + ... + + +@task +def y0(a: typing.List[typing_extensions.Annotated[int, FlyteAnnotation({"foo": {"bar": 1}})]]): + ... + + +@task +def y1(a: typing_extensions.Annotated[typing.List[int], FlyteAnnotation({"foo": {"bar": 1}})]): + ... + + +def test_get_variable_descriptions(): + x_tsk = get_serializable(entity_mapping, serialization_settings, x) + x_input_vars = x_tsk.template.interface.inputs + + a_ann = x_input_vars["a"].type.annotation + assert isinstance(a_ann, TypeAnnotation) + assert a_ann.annotations["foo"] == {"bar": 1} + + b_ann = x_input_vars["b"].type.annotation + assert b_ann is None + + # Annotated simple type within list generic + y0_tsk = get_serializable(entity_mapping, serialization_settings, y0) + y0_input_vars = y0_tsk.template.interface.inputs + y0_a_ann = y0_input_vars["a"].type.collection_type.annotation + assert isinstance(y0_a_ann, TypeAnnotation) + assert y0_a_ann.annotations["foo"] == {"bar": 1} + + # Annotated list generic + y1_tsk = get_serializable(entity_mapping, serialization_settings, y1) + y1_input_vars = y1_tsk.template.interface.inputs + y1_a_ann = y1_input_vars["a"].type.annotation + assert isinstance(y1_a_ann, TypeAnnotation) + assert y1_a_ann.annotations["foo"] == {"bar": 1} diff --git a/tests/flytekit/unit/models/test_types.py b/tests/flytekit/unit/models/test_types.py index 27b5ddf595..82f0b9b519 100644 --- a/tests/flytekit/unit/models/test_types.py +++ b/tests/flytekit/unit/models/test_types.py @@ -2,6 +2,7 @@ from flyteidl.core import types_pb2 from flytekit.models import types as _types +from flytekit.models.annotation import TypeAnnotation from tests.flytekit.common import parameterizers @@ -81,6 +82,16 @@ def test_literal_types(): assert obj == _types.LiteralType.from_flyte_idl(obj.to_flyte_idl()) +def test_annotated_literal_types(): + obj = _types.LiteralType(simple=_types.SimpleType.INTEGER, annotation=TypeAnnotation(annotations={"foo": "bar"})) + assert obj.simple == _types.SimpleType.INTEGER + assert obj.schema is None + assert obj.collection_type is None + assert obj.map_value_type is None + assert obj.annotation.annotations == {"foo": "bar"} + assert obj == _types.LiteralType.from_flyte_idl(obj.to_flyte_idl()) + + @pytest.mark.parametrize("literal_type", parameterizers.LIST_OF_ALL_LITERAL_TYPES) def test_literal_collections(literal_type): obj = _types.LiteralType(collection_type=literal_type) From 2ff05d0d52ffb7558ddddf4ee6d814b04339dd46 Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Tue, 8 Feb 2022 14:02:55 -0800 Subject: [PATCH 079/128] Remove singleton from structured dataset transformer engine (#848) Signed-off-by: Yee Hing Tong Signed-off-by: Kevin Su --- .github/workflows/pythonbuild.yml | 4 +- flytekit/__init__.py | 1 + flytekit/types/structured/basic_dfs.py | 10 ++-- flytekit/types/structured/bigquery.py | 10 ++-- .../types/structured/structured_dataset.py | 47 ++++++++------- .../flytekit-papermill/dev-requirements.in | 3 +- .../flytekit-papermill/dev-requirements.txt | 13 +++-- .../flytekitplugins/spark/sd_transformers.py | 6 +- .../unit/core/test_structured_dataset.py | 57 ++++++++++--------- .../test_structured_dataset_workflow.py | 10 ++-- 10 files changed, 88 insertions(+), 73 deletions(-) diff --git a/.github/workflows/pythonbuild.yml b/.github/workflows/pythonbuild.yml index eda870cff4..56d26badd3 100644 --- a/.github/workflows/pythonbuild.yml +++ b/.github/workflows/pythonbuild.yml @@ -38,9 +38,11 @@ jobs: run: | make setup${{ matrix.spark-version-suffix }} pip freeze + - name: Test FlyteSchema compatibility + run: | + FLYTE_SDK_USE_STRUCTURED_DATASET=FALSE python -m pytest tests/flytekit_compatibility - name: Test with coverage run: | - coverage run -m pytest tests/flytekit_compatibility FLYTE_SDK_USE_STRUCTURED_DATASET=TRUE coverage run -m pytest tests/flytekit/unit - name: Integration Tests with coverage # https://github.com/actions/runner/issues/241#issuecomment-577360161 diff --git a/flytekit/__init__.py b/flytekit/__init__.py index 64c0a26530..6e2e614702 100644 --- a/flytekit/__init__.py +++ b/flytekit/__init__.py @@ -193,6 +193,7 @@ from flytekit.types.structured.structured_dataset import ( StructuredDataset, StructuredDatasetFormat, + StructuredDatasetTransformerEngine, StructuredDatasetType, ) diff --git a/flytekit/types/structured/basic_dfs.py b/flytekit/types/structured/basic_dfs.py index a65fbbeee3..49b2f13ed9 100644 --- a/flytekit/types/structured/basic_dfs.py +++ b/flytekit/types/structured/basic_dfs.py @@ -12,13 +12,13 @@ from flytekit.models.literals import StructuredDatasetMetadata from flytekit.models.types import StructuredDatasetType from flytekit.types.structured.structured_dataset import ( - FLYTE_DATASET_TRANSFORMER, LOCAL, PARQUET, S3, StructuredDataset, StructuredDatasetDecoder, StructuredDatasetEncoder, + StructuredDatasetTransformerEngine, ) T = TypeVar("T") @@ -107,7 +107,7 @@ def decode( for protocol in [LOCAL, S3]: # Should we add GCS - FLYTE_DATASET_TRANSFORMER.register_handler(PandasToParquetEncodingHandler(protocol), default_for_type=True) - FLYTE_DATASET_TRANSFORMER.register_handler(ParquetToPandasDecodingHandler(protocol), default_for_type=True) - FLYTE_DATASET_TRANSFORMER.register_handler(ArrowToParquetEncodingHandler(protocol), default_for_type=True) - FLYTE_DATASET_TRANSFORMER.register_handler(ParquetToArrowDecodingHandler(protocol), default_for_type=True) + StructuredDatasetTransformerEngine.register(PandasToParquetEncodingHandler(protocol), default_for_type=True) + StructuredDatasetTransformerEngine.register(ParquetToPandasDecodingHandler(protocol), default_for_type=True) + StructuredDatasetTransformerEngine.register(ArrowToParquetEncodingHandler(protocol), default_for_type=True) + StructuredDatasetTransformerEngine.register(ParquetToArrowDecodingHandler(protocol), default_for_type=True) diff --git a/flytekit/types/structured/bigquery.py b/flytekit/types/structured/bigquery.py index d9221e4110..923ea06e9e 100644 --- a/flytekit/types/structured/bigquery.py +++ b/flytekit/types/structured/bigquery.py @@ -12,11 +12,11 @@ from flytekit.types.structured.structured_dataset import ( BIGQUERY, DF, - FLYTE_DATASET_TRANSFORMER, StructuredDataset, StructuredDatasetDecoder, StructuredDatasetEncoder, StructuredDatasetMetadata, + StructuredDatasetTransformerEngine, ) @@ -110,7 +110,7 @@ def decode( return pa.Table.from_pandas(_read_from_bq(flyte_value)) -FLYTE_DATASET_TRANSFORMER.register_handler(PandasToBQEncodingHandlers(), default_for_type=False) -FLYTE_DATASET_TRANSFORMER.register_handler(BQToPandasDecodingHandler(), default_for_type=False) -FLYTE_DATASET_TRANSFORMER.register_handler(ArrowToBQEncodingHandlers(), default_for_type=False) -FLYTE_DATASET_TRANSFORMER.register_handler(BQToArrowDecodingHandler(), default_for_type=False) +StructuredDatasetTransformerEngine.register(PandasToBQEncodingHandlers(), default_for_type=False) +StructuredDatasetTransformerEngine.register(BQToPandasDecodingHandler(), default_for_type=False) +StructuredDatasetTransformerEngine.register(ArrowToBQEncodingHandlers(), default_for_type=False) +StructuredDatasetTransformerEngine.register(BQToArrowDecodingHandler(), default_for_type=False) diff --git a/flytekit/types/structured/structured_dataset.py b/flytekit/types/structured/structured_dataset.py index ffefd101dd..819bc012cc 100644 --- a/flytekit/types/structured/structured_dataset.py +++ b/flytekit/types/structured/structured_dataset.py @@ -24,8 +24,7 @@ from flytekit.configuration.sdk import USE_STRUCTURED_DATASET from flytekit.core.context_manager import FlyteContext, FlyteContextManager -from flytekit.core.type_engine import TypeTransformer -from flytekit.extend import TypeEngine +from flytekit.core.type_engine import TypeEngine, TypeTransformer from flytekit.loggers import logger from flytekit.models import literals from flytekit.models import types as type_models @@ -106,7 +105,7 @@ def all(self) -> DF: if self._dataframe_type is None: raise ValueError("No dataframe type set. Use open() to set the local dataframe type you want to use.") ctx = FlyteContextManager.current_context() - return FLYTE_DATASET_TRANSFORMER.open_as( + return flyte_dataset_transformer.open_as( ctx, self.literal, self._dataframe_type, updated_metadata=self.metadata ) @@ -114,7 +113,7 @@ def iter(self) -> Generator[DF, None, None]: if self._dataframe_type is None: raise ValueError("No dataframe type set. Use open() to set the local dataframe type you want to use.") ctx = FlyteContextManager.current_context() - return FLYTE_DATASET_TRANSFORMER.iter_as( + return flyte_dataset_transformer.iter_as( ctx, self.literal, self._dataframe_type, updated_metadata=self.metadata ) @@ -170,7 +169,7 @@ class StructuredDatasetEncoder(ABC): def __init__(self, python_type: Type[T], protocol: str, supported_format: Optional[str] = None): """ Extend this abstract class, implement the encode function, and register your concrete class with the - FLYTE_DATASET_TRANSFORMER defined at this module level in order for the core flytekit type engine to handle + StructuredDatasetTransformerEngine class in order for the core flytekit type engine to handle dataframe libraries. This is the encoding interface, meaning it is used when there is a Python value that the flytekit type engine is trying to convert into a Flyte Literal. For the other way, see the StructuredDatasetEncoder @@ -230,7 +229,7 @@ class StructuredDatasetDecoder(ABC): def __init__(self, python_type: Type[DF], protocol: str, supported_format: Optional[str] = None): """ Extend this abstract class, implement the decode function, and register your concrete class with the - FLYTE_DATASET_TRANSFORMER defined at this module level in order for the core flytekit type engine to handle + StructuredDatasetTransformerEngine class in order for the core flytekit type engine to handle dataframe libraries. This is the decoder interface, meaning it is used when there is a Flyte Literal value, and we have to get a Python value out of it. For the other way, see the StructuredDatasetEncoder @@ -337,7 +336,8 @@ class StructuredDatasetTransformerEngine(TypeTransformer[StructuredDataset]): Handlers = Union[StructuredDatasetEncoder, StructuredDatasetDecoder] - def _finder(self, handler_map, df_type: Type, protocol: str, format: str): + @staticmethod + def _finder(handler_map, df_type: Type, protocol: str, format: str): try: return handler_map[df_type][protocol][format] except KeyError: @@ -352,18 +352,21 @@ def _finder(self, handler_map, df_type: Type, protocol: str, format: str): ... raise ValueError(f"Failed to find a handler for {df_type}, protocol {protocol}, fmt {format}") - def get_encoder(self, df_type: Type, protocol: str, format: str): - return self._finder(self.ENCODERS, df_type, protocol, format) + @classmethod + def get_encoder(cls, df_type: Type, protocol: str, format: str): + return cls._finder(StructuredDatasetTransformerEngine.ENCODERS, df_type, protocol, format) - def get_decoder(self, df_type: Type, protocol: str, format: str): - return self._finder(self.DECODERS, df_type, protocol, format) + @classmethod + def get_decoder(cls, df_type: Type, protocol: str, format: str): + return cls._finder(StructuredDatasetTransformerEngine.DECODERS, df_type, protocol, format) - def _handler_finder(self, h: Handlers) -> Dict[str, Handlers]: + @classmethod + def _handler_finder(cls, h: Handlers) -> Dict[str, Handlers]: # Maybe think about default dict in the future, but is typing as nice? if isinstance(h, StructuredDatasetEncoder): - top_level = self.ENCODERS + top_level = cls.ENCODERS elif isinstance(h, StructuredDatasetDecoder): - top_level = self.DECODERS + top_level = cls.DECODERS else: raise TypeError(f"We don't support this type of handler {h}") if h.python_type not in top_level: @@ -376,7 +379,8 @@ def __init__(self): super().__init__("StructuredDataset Transformer", StructuredDataset) self._type_assertions_enabled = False - def register_handler(self, h: Handlers, default_for_type: Optional[bool] = True, override: Optional[bool] = False): + @classmethod + def register(cls, h: Handlers, default_for_type: Optional[bool] = True, override: Optional[bool] = False): """ Call this with any handler to register it with this dataframe meta-transformer @@ -386,7 +390,7 @@ def register_handler(self, h: Handlers, default_for_type: Optional[bool] = True, logger.info(f"Structured datasets not enabled, not registering handler {h}") return - lowest_level = self._handler_finder(h) + lowest_level = cls._handler_finder(h) if h.supported_format in lowest_level and override is False: raise ValueError(f"Already registered a handler for {(h.python_type, h.protocol, h.supported_format)}") lowest_level[h.supported_format] = h @@ -394,13 +398,14 @@ def register_handler(self, h: Handlers, default_for_type: Optional[bool] = True, if default_for_type: # TODO: Add logging, think about better ux, maybe default False and warn if doesn't exist. - self.DEFAULT_FORMATS[h.python_type] = h.supported_format - self.DEFAULT_PROTOCOLS[h.python_type] = h.protocol + cls.DEFAULT_FORMATS[h.python_type] = h.supported_format + cls.DEFAULT_PROTOCOLS[h.python_type] = h.protocol # Register with the type engine as well # The semantics as of now are such that it doesn't matter which order these transformers are loaded in, as # long as the older Pandas/FlyteSchema transformer do not also specify the override - TypeEngine.register_additional_type(self, h.python_type, override=True) + engine = StructuredDatasetTransformerEngine() + TypeEngine.register_additional_type(engine, h.python_type, override=True) def assert_type(self, t: Type[StructuredDataset], v: typing.Any): return @@ -724,7 +729,7 @@ def guess_python_type(self, literal_type: LiteralType) -> Type[T]: if USE_STRUCTURED_DATASET.get(): logger.debug("Structured dataset module load... using structured datasets!") - FLYTE_DATASET_TRANSFORMER = StructuredDatasetTransformerEngine() - TypeEngine.register(FLYTE_DATASET_TRANSFORMER) + flyte_dataset_transformer = StructuredDatasetTransformerEngine() + TypeEngine.register(flyte_dataset_transformer) else: logger.debug("Structured dataset module load... not using structured datasets") diff --git a/plugins/flytekit-papermill/dev-requirements.in b/plugins/flytekit-papermill/dev-requirements.in index 8bf53e2636..32404e0a2d 100644 --- a/plugins/flytekit-papermill/dev-requirements.in +++ b/plugins/flytekit-papermill/dev-requirements.in @@ -1 +1,2 @@ -flytekitplugins-spark>=0.30.0b4 +git+https://github.com/flyteorg/flytekit@add-sd-make-class-methods#egg=flytekitplugins-spark&subdirectory=plugins/flytekit-spark +# vcs+protocol://repo_url/#egg=pkg&subdirectory=flyte diff --git a/plugins/flytekit-papermill/dev-requirements.txt b/plugins/flytekit-papermill/dev-requirements.txt index d647708f20..ef30545338 100644 --- a/plugins/flytekit-papermill/dev-requirements.txt +++ b/plugins/flytekit-papermill/dev-requirements.txt @@ -40,9 +40,9 @@ docstring-parser==0.13 # via flytekit flyteidl==0.22.0 # via flytekit -flytekit==0.26.0 +flytekit==0.30.0 # via flytekitplugins-spark -flytekitplugins-spark==0.30.0b4 +flytekitplugins-spark @ git+https://github.com/flyteorg/flytekit@add-sd-make-class-methods#subdirectory=plugins/flytekit-spark # via -r dev-requirements.in grpcio==1.43.0 # via flytekit @@ -87,11 +87,11 @@ protobuf==3.19.3 # flytekit py==1.11.0 # via retry -py4j==0.10.9.2 +py4j==0.10.9.3 # via pyspark pyarrow==6.0.1 # via flytekit -pyspark==3.2.0 +pyspark==3.2.1 # via flytekitplugins-spark python-dateutil==2.8.1 # via @@ -123,7 +123,6 @@ retry==0.9.2 six==1.16.0 # via # cookiecutter - # flytekit # grpcio # python-dateutil # responses @@ -134,7 +133,9 @@ statsd==3.3.0 text-unidecode==1.3 # via python-slugify typing-extensions==4.0.1 - # via typing-inspect + # via + # flytekit + # typing-inspect typing-inspect==0.7.1 # via dataclasses-json urllib3==1.26.8 diff --git a/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py b/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py index e0b1c7b41e..2466d3fc13 100644 --- a/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py +++ b/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py @@ -7,11 +7,11 @@ from flytekit.models.literals import StructuredDatasetMetadata from flytekit.models.types import StructuredDatasetType from flytekit.types.structured.structured_dataset import ( - FLYTE_DATASET_TRANSFORMER, PARQUET, StructuredDataset, StructuredDatasetDecoder, StructuredDatasetEncoder, + StructuredDatasetTransformerEngine, ) @@ -45,5 +45,5 @@ def decode( for protocol in ["/", "s3"]: - FLYTE_DATASET_TRANSFORMER.register_handler(SparkToParquetEncodingHandler(protocol), default_for_type=True) - FLYTE_DATASET_TRANSFORMER.register_handler(ParquetToSparkDecodingHandler(protocol), default_for_type=True) + StructuredDatasetTransformerEngine.register(SparkToParquetEncodingHandler(protocol), default_for_type=True) + StructuredDatasetTransformerEngine.register(ParquetToSparkDecodingHandler(protocol), default_for_type=True) diff --git a/tests/flytekit/unit/core/test_structured_dataset.py b/tests/flytekit/unit/core/test_structured_dataset.py index 4e4309e292..a7ef1ea953 100644 --- a/tests/flytekit/unit/core/test_structured_dataset.py +++ b/tests/flytekit/unit/core/test_structured_dataset.py @@ -19,11 +19,11 @@ from flytekit import kwtypes from flytekit.types.structured.structured_dataset import ( - FLYTE_DATASET_TRANSFORMER, PARQUET, StructuredDataset, StructuredDatasetDecoder, StructuredDatasetEncoder, + StructuredDatasetTransformerEngine, convert_schema_type_to_structured_dataset_type, extract_cols_and_format, protocol_prefix, @@ -117,10 +117,10 @@ def test_types_sd(): def test_retrieving(): - assert FLYTE_DATASET_TRANSFORMER.get_encoder(pd.DataFrame, "/", PARQUET) is not None + assert StructuredDatasetTransformerEngine.get_encoder(pd.DataFrame, "/", PARQUET) is not None with pytest.raises(ValueError): # We don't have a default "" format encoder - FLYTE_DATASET_TRANSFORMER.get_encoder(pd.DataFrame, "/", "") + StructuredDatasetTransformerEngine.get_encoder(pd.DataFrame, "/", "") class TempEncoder(StructuredDatasetEncoder): def __init__(self, protocol): @@ -129,15 +129,15 @@ def __init__(self, protocol): def encode(self): ... - FLYTE_DATASET_TRANSFORMER.register_handler(TempEncoder("gs"), default_for_type=False) + StructuredDatasetTransformerEngine.register(TempEncoder("gs"), default_for_type=False) with pytest.raises(ValueError): - FLYTE_DATASET_TRANSFORMER.register_handler(TempEncoder("gs://"), default_for_type=False) + StructuredDatasetTransformerEngine.register(TempEncoder("gs://"), default_for_type=False) class TempEncoder: pass with pytest.raises(TypeError, match="We don't support this type of handler"): - FLYTE_DATASET_TRANSFORMER.register_handler(TempEncoder, default_for_type=False) + StructuredDatasetTransformerEngine.register(TempEncoder, default_for_type=False) def test_to_literal(): @@ -145,7 +145,9 @@ def test_to_literal(): lt = TypeEngine.to_literal_type(pd.DataFrame) df = generate_pandas() - lit = FLYTE_DATASET_TRANSFORMER.to_literal(ctx, df, python_type=pd.DataFrame, expected=lt) + fdt = StructuredDatasetTransformerEngine() + + lit = fdt.to_literal(ctx, df, python_type=pd.DataFrame, expected=lt) assert lit.scalar.structured_dataset.metadata.structured_dataset_type.format == PARQUET assert lit.scalar.structured_dataset.metadata.structured_dataset_type.format == PARQUET @@ -153,16 +155,16 @@ def test_to_literal(): sd_with_literal_and_df._literal_sd = lit with pytest.raises(ValueError, match="Shouldn't have specified both literal"): - FLYTE_DATASET_TRANSFORMER.to_literal(ctx, sd_with_literal_and_df, python_type=StructuredDataset, expected=lt) + fdt.to_literal(ctx, sd_with_literal_and_df, python_type=StructuredDataset, expected=lt) sd_with_nothing = StructuredDataset() with pytest.raises(ValueError, match="If dataframe is not specified"): - FLYTE_DATASET_TRANSFORMER.to_literal(ctx, sd_with_nothing, python_type=StructuredDataset, expected=lt) + fdt.to_literal(ctx, sd_with_nothing, python_type=StructuredDataset, expected=lt) sd_with_uri = StructuredDataset(uri="s3://some/extant/df.parquet") lt = TypeEngine.to_literal_type(Annotated[StructuredDataset, {}, "new-df-format"]) - lit = FLYTE_DATASET_TRANSFORMER.to_literal(ctx, sd_with_uri, python_type=StructuredDataset, expected=lt) + lit = fdt.to_literal(ctx, sd_with_uri, python_type=StructuredDataset, expected=lt) assert lit.scalar.structured_dataset.uri == "s3://some/extant/df.parquet" assert lit.scalar.structured_dataset.metadata.structured_dataset_type.format == "new-df-format" @@ -184,21 +186,22 @@ def encode( ) -> literals.StructuredDataset: return literals.StructuredDataset(uri="") - FLYTE_DATASET_TRANSFORMER.register_handler(TempEncoder("myavro"), default_for_type=True) + StructuredDatasetTransformerEngine.register(TempEncoder("myavro"), default_for_type=True) lt = TypeEngine.to_literal_type(MyDF) assert lt.structured_dataset_type.format == "myavro" ctx = FlyteContextManager.current_context() + fdt = StructuredDatasetTransformerEngine() sd = StructuredDataset(dataframe=42) - l = FLYTE_DATASET_TRANSFORMER.to_literal(ctx, sd, MyDF, lt) + l = fdt.to_literal(ctx, sd, MyDF, lt) # Test that the literal type is filled in even though the encode function above doesn't do it. assert l.scalar.structured_dataset.metadata.structured_dataset_type.format == "myavro" # Test that looking up encoders/decoders falls back to the "" encoder/decoder empty_format_temp_encoder = TempEncoder("") - FLYTE_DATASET_TRANSFORMER.register_handler(empty_format_temp_encoder, default_for_type=False) + StructuredDatasetTransformerEngine.register(empty_format_temp_encoder, default_for_type=False) - res = FLYTE_DATASET_TRANSFORMER.get_encoder(MyDF, "tmpfs", "rando") + res = StructuredDatasetTransformerEngine.get_encoder(MyDF, "tmpfs", "rando") assert res is empty_format_temp_encoder @@ -221,7 +224,7 @@ def decode( ) -> typing.Union[typing.Generator[pd.DataFrame, None, None]]: yield pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]}) - FLYTE_DATASET_TRANSFORMER.register_handler( + StructuredDatasetTransformerEngine.register( MockPandasDecodingHandlers(pd.DataFrame, "tmpfs"), default_for_type=False ) sd = StructuredDataset() @@ -241,7 +244,7 @@ def decode( ) -> pd.DataFrame: pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]}) - FLYTE_DATASET_TRANSFORMER.register_handler( + StructuredDatasetTransformerEngine.register( MockPandasDecodingHandlers(pd.DataFrame, "tmpfs"), default_for_type=False, override=True ) sd = StructuredDataset() @@ -271,24 +274,25 @@ def test_to_python_value_with_incoming_columns(): ctx = FlyteContextManager.current_context() lt = TypeEngine.to_literal_type(original_type) df = generate_pandas() - lit = FLYTE_DATASET_TRANSFORMER.to_literal(ctx, df, python_type=original_type, expected=lt) + fdt = StructuredDatasetTransformerEngine() + lit = fdt.to_literal(ctx, df, python_type=original_type, expected=lt) assert len(lit.scalar.structured_dataset.metadata.structured_dataset_type.columns) == 2 # declare a new type that only has one column # get the dataframe, make sure it has the column that was asked for. subset_sd_type = Annotated[StructuredDataset, kwtypes(age=int)] - sd = FLYTE_DATASET_TRANSFORMER.to_python_value(ctx, lit, subset_sd_type) + sd = fdt.to_python_value(ctx, lit, subset_sd_type) assert sd.metadata.structured_dataset_type.columns[0].name == "age" sub_df = sd.open(pd.DataFrame).all() assert sub_df.shape[1] == 1 # check when columns are not specified, should pull both and add column information. - sd = FLYTE_DATASET_TRANSFORMER.to_python_value(ctx, lit, StructuredDataset) + sd = fdt.to_python_value(ctx, lit, StructuredDataset) assert sd.metadata.structured_dataset_type.columns[0].name == "age" # should also work if subset type is just an annotated pd.DataFrame subset_pd_type = Annotated[pd.DataFrame, kwtypes(age=int)] - sub_df = FLYTE_DATASET_TRANSFORMER.to_python_value(ctx, lit, subset_pd_type) + sub_df = fdt.to_python_value(ctx, lit, subset_pd_type) assert sub_df.shape[1] == 1 @@ -297,13 +301,14 @@ def test_to_python_value_without_incoming_columns(): ctx = FlyteContextManager.current_context() lt = TypeEngine.to_literal_type(pd.DataFrame) df = generate_pandas() - lit = FLYTE_DATASET_TRANSFORMER.to_literal(ctx, df, python_type=pd.DataFrame, expected=lt) + fdt = StructuredDatasetTransformerEngine() + lit = fdt.to_literal(ctx, df, python_type=pd.DataFrame, expected=lt) assert len(lit.scalar.structured_dataset.metadata.structured_dataset_type.columns) == 0 # declare a new type that only has one column # get the dataframe, make sure it has the column that was asked for. subset_sd_type = Annotated[StructuredDataset, kwtypes(age=int)] - sd = FLYTE_DATASET_TRANSFORMER.to_python_value(ctx, lit, subset_sd_type) + sd = fdt.to_python_value(ctx, lit, subset_sd_type) assert sd.metadata.structured_dataset_type.columns[0].name == "age" sub_df = sd.open(pd.DataFrame).all() assert sub_df.shape[1] == 1 @@ -311,14 +316,14 @@ def test_to_python_value_without_incoming_columns(): # check when columns are not specified, should pull both and add column information. # todo: see the todos in the open_as, and iter_as functions in StructuredDatasetTransformerEngine # we have to recreate the literal because the test case above filled in the metadata - lit = FLYTE_DATASET_TRANSFORMER.to_literal(ctx, df, python_type=pd.DataFrame, expected=lt) - sd = FLYTE_DATASET_TRANSFORMER.to_python_value(ctx, lit, StructuredDataset) + lit = fdt.to_literal(ctx, df, python_type=pd.DataFrame, expected=lt) + sd = fdt.to_python_value(ctx, lit, StructuredDataset) assert sd.metadata.structured_dataset_type.columns == [] sub_df = sd.open(pd.DataFrame).all() assert sub_df.shape[1] == 2 # should also work if subset type is just an annotated pd.DataFrame - lit = FLYTE_DATASET_TRANSFORMER.to_literal(ctx, df, python_type=pd.DataFrame, expected=lt) + lit = fdt.to_literal(ctx, df, python_type=pd.DataFrame, expected=lt) subset_pd_type = Annotated[pd.DataFrame, kwtypes(age=int)] - sub_df = FLYTE_DATASET_TRANSFORMER.to_python_value(ctx, lit, subset_pd_type) + sub_df = fdt.to_python_value(ctx, lit, subset_pd_type) assert sub_df.shape[1] == 1 diff --git a/tests/flytekit/unit/types/structured_dataset/test_structured_dataset_workflow.py b/tests/flytekit/unit/types/structured_dataset/test_structured_dataset_workflow.py index fa879ff0d9..d911f971e4 100644 --- a/tests/flytekit/unit/types/structured_dataset/test_structured_dataset_workflow.py +++ b/tests/flytekit/unit/types/structured_dataset/test_structured_dataset_workflow.py @@ -18,13 +18,13 @@ from flytekit.types.structured.structured_dataset import ( BIGQUERY, DF, - FLYTE_DATASET_TRANSFORMER, LOCAL, PARQUET, S3, StructuredDataset, StructuredDatasetDecoder, StructuredDatasetEncoder, + StructuredDatasetTransformerEngine, ) PANDAS_PATH = FlyteContextManager.current_context().file_access.get_random_local_directory() @@ -58,8 +58,8 @@ def decode( return pd_df -FLYTE_DATASET_TRANSFORMER.register_handler(MockBQEncodingHandlers(pd.DataFrame, BIGQUERY), False, True) -FLYTE_DATASET_TRANSFORMER.register_handler(MockBQDecodingHandlers(pd.DataFrame, BIGQUERY), False, True) +StructuredDatasetTransformerEngine.register(MockBQEncodingHandlers(pd.DataFrame, BIGQUERY), False, True) +StructuredDatasetTransformerEngine.register(MockBQDecodingHandlers(pd.DataFrame, BIGQUERY), False, True) class NumpyEncodingHandlers(StructuredDatasetEncoder): @@ -95,8 +95,8 @@ def decode( for protocol in [LOCAL, S3]: - FLYTE_DATASET_TRANSFORMER.register_handler(NumpyEncodingHandlers(np.ndarray, protocol, PARQUET)) - FLYTE_DATASET_TRANSFORMER.register_handler(NumpyDecodingHandlers(np.ndarray, protocol, PARQUET)) + StructuredDatasetTransformerEngine.register(NumpyEncodingHandlers(np.ndarray, protocol, PARQUET)) + StructuredDatasetTransformerEngine.register(NumpyDecodingHandlers(np.ndarray, protocol, PARQUET)) @task From a33c108e353ddcf8edc5f562095c6ce37077caa8 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Thu, 10 Feb 2022 05:48:34 +0800 Subject: [PATCH 080/128] Assign input and output to FlyteWorkflowExecution (#842) Signed-off-by: Kevin Su Signed-off-by: maximsmol --- flytekit/remote/remote.py | 3 ++- tests/flytekit/integration/remote/test_remote.py | 4 ++++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/flytekit/remote/remote.py b/flytekit/remote/remote.py index 44d5e8617c..93adc04db8 100644 --- a/flytekit/remote/remote.py +++ b/flytekit/remote/remote.py @@ -460,7 +460,7 @@ def fetch_workflow_execution( """ if name is None: raise user_exceptions.FlyteAssertion("the 'name' argument must be specified.") - return FlyteWorkflowExecution.promote_from_model( + execution = FlyteWorkflowExecution.promote_from_model( self.client.get_execution( WorkflowExecutionIdentifier( project or self.default_project, @@ -469,6 +469,7 @@ def fetch_workflow_execution( ) ) ) + return self.sync_workflow_execution(execution) ###################### # Listing Entities # diff --git a/tests/flytekit/integration/remote/test_remote.py b/tests/flytekit/integration/remote/test_remote.py index 829137c928..ba105d386b 100644 --- a/tests/flytekit/integration/remote/test_remote.py +++ b/tests/flytekit/integration/remote/test_remote.py @@ -192,6 +192,10 @@ def test_execute_python_workflow_and_launch_plan(flyteclient, flyte_workflows_re assert execution.outputs["o0"] == 16 assert execution.outputs["o1"] == "foobarworld" + flyte_workflow_execution = remote.fetch_workflow_execution(name=execution.id.name) + assert execution.inputs == flyte_workflow_execution.inputs + assert execution.outputs == flyte_workflow_execution.outputs + def test_fetch_execute_launch_plan_list_of_floats(flyteclient, flyte_workflows_register): remote = FlyteRemote.from_config(PROJECT, "development") From 0af55a2a96bd2c3265cd52a4de42e72aca8adcb6 Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Fri, 11 Feb 2022 18:53:00 -0800 Subject: [PATCH 081/128] Add reference entities to FlyteTask and FlyteLaunchPlan (#850) Signed-off-by: Yee Hing Tong Signed-off-by: maximsmol --- flytekit/remote/launch_plan.py | 45 +++++++++-- flytekit/remote/task.py | 37 ++++++++- flytekit/remote/workflow.py | 1 - plugins/flytekit-greatexpectations/setup.py | 2 +- tests/flytekit/unit/remote/test_calling.py | 85 +++++++++++++++++++++ 5 files changed, 159 insertions(+), 11 deletions(-) create mode 100644 tests/flytekit/unit/remote/test_calling.py diff --git a/flytekit/remote/launch_plan.py b/flytekit/remote/launch_plan.py index b5dddea03f..3bf845fcab 100644 --- a/flytekit/remote/launch_plan.py +++ b/flytekit/remote/launch_plan.py @@ -1,7 +1,9 @@ from typing import Optional from flytekit.core.interface import Interface +from flytekit.core.launch_plan import ReferenceLaunchPlan from flytekit.core.type_engine import TypeEngine +from flytekit.loggers import remote_logger as logger from flytekit.models import interface as _interface_models from flytekit.models import launch_plan as _launch_plan_models from flytekit.models.core import identifier as id_models @@ -18,8 +20,41 @@ def __init__(self, id, *args, **kwargs): # The interface is not set explicitly unless fetched in an engine context self._interface = None - self._python_interface = None + self._reference_entity = None + + def __call__(self, *args, **kwargs): + if self.reference_entity is None: + logger.warning( + f"FlyteLaunchPlan {self} is not callable, most likely because flytekit could not " + f"guess the python interface. The workflow calling this launch plan may not behave correctly." + ) + return + return self.reference_entity(*args, **kwargs) + + # TODO: Refactor behind mixin + @property + def reference_entity(self) -> Optional[ReferenceLaunchPlan]: + if self._reference_entity is None: + if self.guessed_python_interface is None: + try: + self.guessed_python_interface = Interface( + TypeEngine.guess_python_types(self.interface.inputs), + TypeEngine.guess_python_types(self.interface.outputs), + ) + except Exception as e: + logger.warning(f"Error backing out interface {e}, Flyte interface {self.interface}") + return None + + self._reference_entity = ReferenceLaunchPlan( + self.id.project, + self.id.domain, + self.id.name, + self.id.version, + inputs=self.guessed_python_interface.inputs, + outputs=self.guessed_python_interface.outputs, + ) + return self._reference_entity @classmethod def promote_from_model( @@ -37,12 +72,6 @@ def promote_from_model( raw_output_data_config=model.raw_output_data_config, ) - if lp.interface is not None: - lp.guessed_python_interface = Interface( - inputs=TypeEngine.guess_python_types(lp.interface.inputs), - outputs=TypeEngine.guess_python_types(lp.interface.outputs), - ) - return lp @property @@ -65,7 +94,7 @@ def workflow_id(self) -> id_models.Identifier: return self._workflow_id @property - def interface(self) -> _interface.TypedInterface: + def interface(self) -> Optional[_interface.TypedInterface]: """ The interface is not technically part of the admin.LaunchPlanSpec in the IDL, however the workflow ID is, and from the workflow ID, fetch will fill in the interface. This is nice because then you can __call__ the= diff --git a/flytekit/remote/task.py b/flytekit/remote/task.py index 1ff99549d6..34b4f1d9c2 100644 --- a/flytekit/remote/task.py +++ b/flytekit/remote/task.py @@ -2,8 +2,9 @@ from flytekit.core import hash as _hash_mixin from flytekit.core.interface import Interface +from flytekit.core.task import ReferenceTask from flytekit.core.type_engine import TypeEngine -from flytekit.loggers import logger +from flytekit.loggers import remote_logger as logger from flytekit.models import task as _task_model from flytekit.models.core import identifier as _identifier_model from flytekit.remote import interface as _interfaces @@ -24,6 +25,40 @@ def __init__(self, id, type, metadata, interface, custom, container=None, task_t config=config, ) self._python_interface = None + self._reference_entity = None + + def __call__(self, *args, **kwargs): + if self.reference_entity is None: + logger.warning( + f"FlyteTask {self} is not callable, most likely because flytekit could not " + f"guess the python interface. The workflow calling this task may not behave correctly" + ) + return + return self.reference_entity(*args, **kwargs) + + # TODO: Refactor behind mixin + @property + def reference_entity(self) -> Optional[ReferenceTask]: + if self._reference_entity is None: + if self.guessed_python_interface is None: + try: + self.guessed_python_interface = Interface( + TypeEngine.guess_python_types(self.interface.inputs), + TypeEngine.guess_python_types(self.interface.outputs), + ) + except Exception as e: + logger.warning(f"Error backing out interface {e}, Flyte interface {self.interface}") + return None + + self._reference_entity = ReferenceTask( + self.id.project, + self.id.domain, + self.id.name, + self.id.version, + inputs=self.guessed_python_interface.inputs, + outputs=self.guessed_python_interface.outputs, + ) + return self._reference_entity @property def interface(self) -> _interfaces.TypedInterface: diff --git a/flytekit/remote/workflow.py b/flytekit/remote/workflow.py index 010a703187..14ef6c91bf 100644 --- a/flytekit/remote/workflow.py +++ b/flytekit/remote/workflow.py @@ -57,7 +57,6 @@ def __init__( self._tasks = tasks self._launch_plans = launch_plans self._compiled_closure = compiled_closure - self._node_map = None @property diff --git a/plugins/flytekit-greatexpectations/setup.py b/plugins/flytekit-greatexpectations/setup.py index 3541415664..ac7d9481ef 100644 --- a/plugins/flytekit-greatexpectations/setup.py +++ b/plugins/flytekit-greatexpectations/setup.py @@ -4,7 +4,7 @@ microlib_name = f"flytekitplugins-{PLUGIN_NAME}" -plugin_requires = ["flytekit>=0.21.0,<1.0.0", "great-expectations>=0.13.30", "sqlalchemy>=1.4.23"] +plugin_requires = ["flytekit>=0.21.0,<1.0.0", "great-expectations>=0.13.30,<0.14.6", "sqlalchemy>=1.4.23"] __version__ = "0.0.0+develop" diff --git a/tests/flytekit/unit/remote/test_calling.py b/tests/flytekit/unit/remote/test_calling.py new file mode 100644 index 0000000000..97a1a001dc --- /dev/null +++ b/tests/flytekit/unit/remote/test_calling.py @@ -0,0 +1,85 @@ +import typing +from collections import OrderedDict + +import pytest + +from flytekit.core import context_manager +from flytekit.core.context_manager import Image, ImageConfig +from flytekit.core.launch_plan import LaunchPlan +from flytekit.core.reference_entity import ReferenceSpec +from flytekit.core.task import task +from flytekit.core.workflow import workflow +from flytekit.remote import FlyteLaunchPlan, FlyteTask +from flytekit.remote.interface import TypedInterface +from flytekit.tools.translator import gather_dependent_entities, get_serializable + +default_img = Image(name="default", fqn="test", tag="tag") +serialization_settings = context_manager.SerializationSettings( + project="project", + domain="domain", + version="version", + env=None, + image_config=ImageConfig(default_image=default_img, images=[default_img]), +) + + +@task +def t1(a: int) -> int: + return a + 2 + + +@task +def t2(a: int, b: str) -> str: + return b + str(a) + + +@workflow +def sub_wf(a: int, b: str) -> (int, str): + x = t1(a=a) + d = t2(a=x, b=b) + return x, d + + +serialized = OrderedDict() +t1_spec = get_serializable(serialized, serialization_settings, t1) +ft = FlyteTask.promote_from_model(t1_spec.template) + + +def test_fetched_task(): + @workflow + def wf(a: int) -> int: + return ft(a=a) + + # Should not work unless mocked out. + with pytest.raises(Exception, match="cannot be run locally"): + wf(a=3) + + # Should have one reference entity + serialized = OrderedDict() + get_serializable(serialized, serialization_settings, wf) + vals = [v for v in serialized.values()] + refs = [f for f in filter(lambda x: isinstance(x, ReferenceSpec), vals)] + assert len(refs) == 1 + + +def test_calling_lp(): + sub_wf_lp = LaunchPlan.get_or_create(sub_wf) + serialized = OrderedDict() + lp_model = get_serializable(serialized, serialization_settings, sub_wf_lp) + task_templates, wf_specs, lp_specs = gather_dependent_entities(serialized) + for wf_id, spec in wf_specs.items(): + break + + remote_lp = FlyteLaunchPlan.promote_from_model(lp_model.id, lp_model.spec) + # To pretend that we've fetched this launch plan from Admin, also fill in the Flyte interface, which isn't + # part of the IDL object but is something FlyteRemote does + remote_lp._interface = TypedInterface.promote_from_model(spec.template.interface) + serialized = OrderedDict() + + @workflow + def wf2(a: int) -> typing.Tuple[int, str]: + return remote_lp(a=a, b="hello") + + wf_spec = get_serializable(serialized, serialization_settings, wf2) + print(wf_spec.template.nodes[0].workflow_node.launchplan_ref) + assert wf_spec.template.nodes[0].workflow_node.launchplan_ref == lp_model.id From d8f2fb04498e575951958a8adaf483121e6c3b0f Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Mon, 14 Feb 2022 09:07:57 -0800 Subject: [PATCH 082/128] Fix fast registration error (#851) Signed-off-by: Yee Hing Tong Signed-off-by: Samhita Alla Signed-off-by: maximsmol --- flytekit/tools/fast_registration.py | 3 +-- .../great_expectations/schema.py | 25 ++++++++----------- .../great_expectations/task.py | 24 ++++++++---------- .../requirements.txt | 16 ++++++------ 4 files changed, 30 insertions(+), 38 deletions(-) diff --git a/flytekit/tools/fast_registration.py b/flytekit/tools/fast_registration.py index f9b9fc6665..ece363405c 100644 --- a/flytekit/tools/fast_registration.py +++ b/flytekit/tools/fast_registration.py @@ -100,8 +100,7 @@ def download_distribution(additional_distribution: str, destination: str): """ file_access.get_data(additional_distribution, destination) tarfile_name = _os.path.basename(additional_distribution) - file_suffix = _Path(tarfile_name).suffixes - if len(file_suffix) != 2 or file_suffix[0] != ".tar" or file_suffix[1] != ".gz": + if not tarfile_name.endswith(".tar.gz"): raise ValueError("Unrecognized additional distribution format for {}".format(additional_distribution)) # This will overwrite the existing user flyte workflow code in the current working code dir. diff --git a/plugins/flytekit-greatexpectations/flytekitplugins/great_expectations/schema.py b/plugins/flytekit-greatexpectations/flytekitplugins/great_expectations/schema.py index 2dafc26c53..cd6f321dd2 100644 --- a/plugins/flytekit-greatexpectations/flytekitplugins/great_expectations/schema.py +++ b/plugins/flytekit-greatexpectations/flytekitplugins/great_expectations/schema.py @@ -286,27 +286,14 @@ def to_python_value( } ) - checkpoint_config = { - "class_name": "SimpleCheckpoint", - "validations": [ - { - "batch_request": final_batch_request, - "expectation_suite_name": ge_conf.expectation_suite_name, - } - ], - } - if ge_conf.checkpoint_params: checkpoint = SimpleCheckpoint( f"_tmp_checkpoint_{ge_conf.expectation_suite_name}", context, - **checkpoint_config, **ge_conf.checkpoint_params, ) else: - checkpoint = SimpleCheckpoint( - f"_tmp_checkpoint_{ge_conf.expectation_suite_name}", context, **checkpoint_config - ) + checkpoint = SimpleCheckpoint(f"_tmp_checkpoint_{ge_conf.expectation_suite_name}", context) # identify every run uniquely run_id = RunIdentifier( @@ -316,7 +303,15 @@ def to_python_value( } ) - checkpoint_result = checkpoint.run(run_id=run_id) + checkpoint_result = checkpoint.run( + run_id=run_id, + validations=[ + { + "batch_request": final_batch_request, + "expectation_suite_name": ge_conf.expectation_suite_name, + } + ], + ) final_result = convert_to_json_serializable(checkpoint_result.list_validation_results())[0] result_string = "" diff --git a/plugins/flytekit-greatexpectations/flytekitplugins/great_expectations/task.py b/plugins/flytekit-greatexpectations/flytekitplugins/great_expectations/task.py index c395579f51..0ba60217f3 100644 --- a/plugins/flytekit-greatexpectations/flytekitplugins/great_expectations/task.py +++ b/plugins/flytekit-greatexpectations/flytekitplugins/great_expectations/task.py @@ -231,26 +231,16 @@ def execute(self, **kwargs) -> Any: } ) - checkpoint_config = { - "class_name": "SimpleCheckpoint", - "validations": [ - { - "batch_request": final_batch_request, - "expectation_suite_name": self._expectation_suite_name, - } - ], - } - if self._checkpoint_params: checkpoint = SimpleCheckpoint( f"_tmp_checkpoint_{self._expectation_suite_name}", context, - **checkpoint_config, **self._checkpoint_params, ) else: checkpoint = SimpleCheckpoint( - f"_tmp_checkpoint_{self._expectation_suite_name}", context, **checkpoint_config + f"_tmp_checkpoint_{self._expectation_suite_name}", + context, ) # identify every run uniquely @@ -261,7 +251,15 @@ def execute(self, **kwargs) -> Any: } ) - checkpoint_result = checkpoint.run(run_id=run_id) + checkpoint_result = checkpoint.run( + run_id=run_id, + validations=[ + { + "batch_request": final_batch_request, + "expectation_suite_name": self._expectation_suite_name, + } + ], + ) final_result = convert_to_json_serializable(checkpoint_result.list_validation_results())[0] result_string = "" diff --git a/plugins/flytekit-greatexpectations/requirements.txt b/plugins/flytekit-greatexpectations/requirements.txt index 624f3dc633..7aa0a867dd 100644 --- a/plugins/flytekit-greatexpectations/requirements.txt +++ b/plugins/flytekit-greatexpectations/requirements.txt @@ -8,6 +8,10 @@ # via -r requirements.in altair==4.2.0 # via great-expectations +appnope==0.1.2 + # via + # ipykernel + # ipython argon2-cffi==21.3.0 # via notebook argon2-cffi-bindings==21.2.0 @@ -29,9 +33,7 @@ bleach==4.1.0 certifi==2021.10.8 # via requests cffi==1.15.0 - # via - # argon2-cffi-bindings - # cryptography + # via argon2-cffi-bindings chardet==4.0.0 # via binaryornot charset-normalizer==2.0.10 @@ -50,8 +52,6 @@ cookiecutter==1.7.3 # via flytekit croniter==1.2.0 # via flytekit -cryptography==36.0.1 - # via secretstorage dataclasses-json==0.5.6 # via flytekit debugpy==1.5.1 @@ -285,10 +285,10 @@ responses==0.17.0 # via flytekit retry==0.9.2 # via flytekit -ruamel.yaml==0.17.17 +ruamel-yaml==0.17.17 # via great-expectations -ruamel.yaml.clib==0.2.6 - # via ruamel.yaml +ruamel-yaml-clib==0.2.6 + # via ruamel-yaml scipy==1.7.3 # via great-expectations send2trash==1.8.0 From 7442893ea572b6dd582304441ace233e93e3b9db Mon Sep 17 00:00:00 2001 From: ggydush-fn <69013027+ggydush-fn@users.noreply.github.com> Date: Tue, 15 Feb 2022 16:48:38 -0500 Subject: [PATCH 083/128] Add support for local execute in pod task (#852) Signed-off-by: ggydush-fn Signed-off-by: maximsmol --- plugins/flytekit-k8s-pod/flytekitplugins/pod/task.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/plugins/flytekit-k8s-pod/flytekitplugins/pod/task.py b/plugins/flytekit-k8s-pod/flytekitplugins/pod/task.py index a4872bb321..e9d7c93d95 100644 --- a/plugins/flytekit-k8s-pod/flytekitplugins/pod/task.py +++ b/plugins/flytekit-k8s-pod/flytekitplugins/pod/task.py @@ -7,6 +7,7 @@ from flytekit import FlyteContext, PythonFunctionTask from flytekit.exceptions import user as _user_exceptions from flytekit.extend import Promise, SerializationSettings, TaskPlugins +from flytekit.loggers import logger from flytekit.models import task as _task_models _PRIMARY_CONTAINER_NAME_FIELD = "primary_container_name" @@ -120,7 +121,10 @@ def get_config(self, settings: SerializationSettings) -> Dict[str, str]: return {_PRIMARY_CONTAINER_NAME_FIELD: self.task_config.primary_container_name} def local_execute(self, ctx: FlyteContext, **kwargs) -> Union[Tuple[Promise], Promise, None]: - raise _user_exceptions.FlyteUserException("Local execute is not currently supported for pod tasks") + logger.warning( + "Running pod task locally. Local environment may not match pod environment which may cause issues." + ) + return super().local_execute(ctx=ctx, **kwargs) TaskPlugins.register_pythontask_plugin(Pod, PodFunctionTask) From 32186be449716a3ad84d14d1071968eaba46a864 Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Wed, 16 Feb 2022 13:08:31 -0800 Subject: [PATCH 084/128] Add anonymous retry (#854) Signed-off-by: Yee Hing Tong Signed-off-by: maximsmol --- flytekit/extras/persistence/s3_awscli.py | 22 ++++++++++++---- .../flytekitplugins/fsspec/persist.py | 26 ++++++++++++++++--- .../unit/extras/persistence/test_s3_awscli.py | 2 +- 3 files changed, 41 insertions(+), 9 deletions(-) diff --git a/flytekit/extras/persistence/s3_awscli.py b/flytekit/extras/persistence/s3_awscli.py index 3b24fef94b..64e09e219c 100644 --- a/flytekit/extras/persistence/s3_awscli.py +++ b/flytekit/extras/persistence/s3_awscli.py @@ -1,4 +1,3 @@ -import logging import os as _os import re as _re import string as _string @@ -10,8 +9,11 @@ from flytekit.configuration import aws from flytekit.core.data_persistence import DataPersistence, DataPersistencePlugins from flytekit.exceptions.user import FlyteUserException +from flytekit.loggers import logger from flytekit.tools import subprocess +S3_ANONYMOUS_FLAG = "--no-sign-request" + def _update_cmd_config_and_execute(cmd: List[str]): env = _os.environ.copy() @@ -32,16 +34,26 @@ def _update_cmd_config_and_execute(cmd: List[str]): retry = 0 while True: try: - return subprocess.check_call(cmd, env=env) + try: + return subprocess.check_call(cmd, env=env) + except Exception as e: + if retry > 0: + logger.info(f"AWS command failed with error {e}, command: {cmd}, retry {retry}") + + logger.debug(f"Appending anonymous flag and retrying command {cmd}") + anonymous_cmd = cmd[:] # strings only, so this is deep enough + anonymous_cmd.insert(1, S3_ANONYMOUS_FLAG) + return subprocess.check_call(anonymous_cmd, env=env) + except Exception as e: - logging.error(f"Exception when trying to execute {cmd}, reason: {str(e)}") + logger.error(f"Exception when trying to execute {cmd}, reason: {str(e)}") retry += 1 if retry > aws.RETRIES.get(): raise secs = aws.BACKOFF_SECONDS.get() - logging.info(f"Sleeping before retrying again, after {secs} seconds") + logger.info(f"Sleeping before retrying again, after {secs} seconds") time.sleep(secs) - logging.info("Retrying again") + logger.info("Retrying again") def _extra_args(extra_args: Dict[str, str]) -> List[str]: diff --git a/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/persist.py b/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/persist.py index d2ea879ce0..68bc92b493 100644 --- a/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/persist.py +++ b/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/persist.py @@ -67,14 +67,34 @@ def recursive_paths(f: str, t: str) -> typing.Tuple[str, str]: return f, t def exists(self, path: str) -> bool: - fs = self._get_filesystem(path) - return fs.exists(path) + try: + fs = self._get_filesystem(path) + return fs.exists(path) + except OSError as oe: + logger.debug(f"Error in exists checking {path} {oe}") + protocol = FSSpecPersistence._get_protocol(path) + if protocol == "s3": + logger.debug("S3 source detected, attempting anonymous S3 exists check") + kwargs = s3_setup_args() + anonymous_fs = fsspec.filesystem(protocol, anon=True, **kwargs) # type: ignore + return anonymous_fs.exists(path) + raise oe def get(self, from_path: str, to_path: str, recursive: bool = False): fs = self._get_filesystem(from_path) if recursive: from_path, to_path = self.recursive_paths(from_path, to_path) - return fs.get(from_path, to_path, recursive=recursive) + try: + return fs.get(from_path, to_path, recursive=recursive) + except OSError as oe: + logger.debug(f"Error in getting {from_path} to {to_path} rec {recursive} {oe}") + protocol = FSSpecPersistence._get_protocol(from_path) + if protocol == "s3": + logger.debug("S3 source detected, attempting anonymous S3 access") + kwargs = s3_setup_args() + anonymous_fs = fsspec.filesystem(protocol, anon=True, **kwargs) # type: ignore + return anonymous_fs.get(from_path, to_path, recursive=recursive) + raise oe def put(self, from_path: str, to_path: str, recursive: bool = False): fs = self._get_filesystem(to_path) diff --git a/tests/flytekit/unit/extras/persistence/test_s3_awscli.py b/tests/flytekit/unit/extras/persistence/test_s3_awscli.py index 78f7d67b88..bcf1fd3495 100644 --- a/tests/flytekit/unit/extras/persistence/test_s3_awscli.py +++ b/tests/flytekit/unit/extras/persistence/test_s3_awscli.py @@ -25,7 +25,7 @@ def test_retries(mock_subprocess, mock_delay, mock_check): proxy = S3Persistence() assert proxy.exists("s3://test/fdsa/fdsa") is False - assert mock_subprocess.check_call.call_count == 4 + assert mock_subprocess.check_call.call_count == 8 def test_extra_args(): From 813e56e83c130723386e7755dba9559c40449251 Mon Sep 17 00:00:00 2001 From: maximsmol Date: Wed, 2 Mar 2022 11:44:08 -0800 Subject: [PATCH 085/128] fix: formatting, linting, typing_extensions --- flytekit/core/type_engine.py | 9 +++++---- tests/flytekit/unit/core/test_type_engine.py | 3 +-- tests/flytekit/unit/core/test_type_hints.py | 1 - 3 files changed, 6 insertions(+), 7 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index c8e2ce354e..6076e86537 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -8,9 +8,10 @@ import mimetypes import typing from abc import ABC, abstractmethod -from re import L from typing import NamedTuple, Optional, Type, cast +from typing_extensions import get_args as _get_args + try: from typing import Annotated, get_args, get_origin except ImportError: @@ -185,7 +186,7 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: if type(res) != self._type: raise TypeTransformerFailedError(f"Cannot convert literal {lv} to {self._type}") return res - except AttributeError as e: + except AttributeError: # Assume that this is because a property on `lv` was None raise TypeTransformerFailedError(f"Cannot convert literal {lv}") @@ -503,7 +504,7 @@ def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], exp struct = Struct() try: struct.update(_MessageToDict(python_val)) - except: + except Exception: raise TypeTransformerFailedError("Failed to convert to generic protobuf struct") return Literal(scalar=Scalar(generic=struct)) @@ -866,7 +867,7 @@ def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], exp res_type = _add_tag_to_type(trans.get_literal_type(t), trans.name) if found_res: # Should really never happen, sanity check - raise TypeError(f"Ambiguous choice of variant for union type") + raise TypeError("Ambiguous choice of variant for union type") found_res = True except TypeTransformerFailedError as e: logger.debug(f"Failed to convert from {python_val} to {t}", e) diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 7de44b9935..e7280f47cb 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -20,7 +20,6 @@ import flytekit.common.exceptions.user as user_exceptions from flytekit import kwtypes -from flytekit.common.exceptions import user as user_exceptions from flytekit.common.types import primitives from flytekit.core.annotation import FlyteAnnotation from flytekit.core.context_manager import FlyteContext, FlyteContextManager @@ -374,7 +373,7 @@ def test_guessing_basic(): lt = model_types.LiteralType(simple=model_types.SimpleType.NONE) pt = TypeEngine.guess_python_type(lt) - assert pt is type(None) + assert pt is type(None) # noqa: E721 lt = model_types.LiteralType( blob=BlobType( diff --git a/tests/flytekit/unit/core/test_type_hints.py b/tests/flytekit/unit/core/test_type_hints.py index 40442c72d3..34511fa9be 100644 --- a/tests/flytekit/unit/core/test_type_hints.py +++ b/tests/flytekit/unit/core/test_type_hints.py @@ -19,7 +19,6 @@ import flytekit from flytekit import ContainerTask, Secret, SQLTask, dynamic, kwtypes, map_task -from flytekit.common.translator import get_serializable from flytekit.common.types import primitives from flytekit.core import context_manager, launch_plan, promise from flytekit.core.condition import conditional From abbb15c8e642945e0b4e20abdcec2ef5a55d683b Mon Sep 17 00:00:00 2001 From: maximsmol Date: Wed, 2 Mar 2022 11:48:25 -0800 Subject: [PATCH 086/128] fix: do not use SDK types Signed-off-by: maximsmol --- tests/flytekit/unit/core/test_type_engine.py | 9 ++++----- tests/flytekit/unit/core/test_type_hints.py | 5 ++--- 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index e7280f47cb..19bfc29691 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -20,7 +20,6 @@ import flytekit.common.exceptions.user as user_exceptions from flytekit import kwtypes -from flytekit.common.types import primitives from flytekit.core.annotation import FlyteAnnotation from flytekit.core.context_manager import FlyteContext, FlyteContextManager from flytekit.core.type_engine import ( @@ -785,7 +784,7 @@ def test_union_from_unambiguous_literal(): assert union_type_tags_unique(lt) ctx = FlyteContextManager.current_context() - lv = TypeEngine.to_literal(ctx, 3, int, primitives.Integer.to_flyte_literal_type()) + lv = TypeEngine.to_literal(ctx, 3, int, LiteralType(simple=SimpleType.INTEGER)) assert lv.scalar.primitive.integer == 3 v = TypeEngine.to_python_value(ctx, lv, pt) @@ -806,7 +805,7 @@ def __eq__(self, other): SimpleTransformer( "MyInt", MyInt, - primitives.Integer.to_flyte_literal_type(), + LiteralType(simple=SimpleType.INTEGER), lambda x: Literal(scalar=Scalar(primitive=Primitive(integer=x.val))), lambda x: MyInt(x.scalar.primitive.integer), ) @@ -833,7 +832,7 @@ def __eq__(self, other): assert lv.scalar.union.value.scalar.primitive.integer == 10 assert v == MyInt(10) - lv = TypeEngine.to_literal(ctx, 4, int, primitives.Integer.to_flyte_literal_type()) + lv = TypeEngine.to_literal(ctx, 4, int, LiteralType(simple=SimpleType.INTEGER)) assert lv.scalar.primitive.integer == 4 try: TypeEngine.to_python_value(ctx, lv, pt) @@ -859,7 +858,7 @@ def __init__(self): super().__init__("UnsignedInt", UnsignedInt) def get_literal_type(self, t: typing.Type[T]) -> LiteralType: - return primitives.Integer.to_flyte_literal_type() + return LiteralType(simple=SimpleType.INTEGER) def to_literal( self, ctx: FlyteContext, python_val: T, python_type: typing.Type[T], expected: LiteralType diff --git a/tests/flytekit/unit/core/test_type_hints.py b/tests/flytekit/unit/core/test_type_hints.py index 34511fa9be..67247f53ee 100644 --- a/tests/flytekit/unit/core/test_type_hints.py +++ b/tests/flytekit/unit/core/test_type_hints.py @@ -19,7 +19,6 @@ import flytekit from flytekit import ContainerTask, Secret, SQLTask, dynamic, kwtypes, map_task -from flytekit.common.types import primitives from flytekit.core import context_manager, launch_plan, promise from flytekit.core.condition import conditional from flytekit.core.context_manager import ExecutionState, FastSerializationSettings, Image, ImageConfig @@ -1713,7 +1712,7 @@ def __eq__(self, other): SimpleTransformer( "MyInt", MyInt, - primitives.Integer.to_flyte_literal_type(), + LiteralType(simple=SimpleType.INTEGER), lambda x: _literal_models.Literal( scalar=_literal_models.Scalar(primitive=_literal_models.Primitive(integer=x.val)) ), @@ -1753,7 +1752,7 @@ def __eq__(self, other): SimpleTransformer( "MyInt", MyInt, - primitives.Integer.to_flyte_literal_type(), + LiteralType(simple=SimpleType.INTEGER), lambda x: _literal_models.Literal( scalar=_literal_models.Scalar(primitive=_literal_models.Primitive(integer=x.val)) ), From 960b1613ea2a4e92011a54b01bdf4dd00ea5caf8 Mon Sep 17 00:00:00 2001 From: maximsmol Date: Wed, 2 Mar 2022 11:54:42 -0800 Subject: [PATCH 087/128] fix: update test comment Signed-off-by: maximsmol --- tests/flytekit/unit/core/test_type_engine.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 19bfc29691..e511cceb6f 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -905,7 +905,11 @@ def test_union_of_lists(): structure=TypeStructure(tag="Typed List"), ), ] - assert not union_type_tags_unique(lt) # tags are deliberately NOT unique + # Tags are deliberately NOT unique beacuse they are not required to encode the deep type structure, + # only the top-level type transformer choice + # + # The stored typed will be used to differentiate union variants and must produce a unique choice. + assert not union_type_tags_unique(lt) ctx = FlyteContextManager.current_context() lv = TypeEngine.to_literal(ctx, ["hello", "world"], pt, lt) From 1368024bcc636c948d04e97faa5a223d88cd0ed1 Mon Sep 17 00:00:00 2001 From: maximsmol Date: Wed, 2 Mar 2022 17:18:31 -0800 Subject: [PATCH 088/128] fix: also check literal type castability when tags match Signed-off-by: maximsmol --- flytekit/core/type_engine.py | 98 ++++++++++++++++++++++++++++++++++-- flytekit/types/file/file.py | 6 ++- 2 files changed, 99 insertions(+), 5 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 6076e86537..7a405df6e3 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -824,8 +824,13 @@ def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], exp return Literal(collection=LiteralCollection(literals=lit_list)) def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[T]) -> typing.List[T]: + try: + lits = lv.collection.literals + except AttributeError: + raise TypeTransformerFailedError() + st = self.get_sub_type(expected_python_type) - return [TypeEngine.to_python_value(ctx, x, st) for x in lv.collection.literals] + return [TypeEngine.to_python_value(ctx, x, st) for x in lits] def guess_python_type(self, literal_type: LiteralType) -> Type[list]: if literal_type.collection_type: @@ -839,6 +844,85 @@ def _add_tag_to_type(x: LiteralType, tag: str) -> LiteralType: return x +def _type_essence(x: LiteralType) -> LiteralType: + if x.metadata is not None or x.structure is not None or x.annotation is not None: + x = LiteralType.from_flyte_idl(x.to_flyte_idl()) + x._metadata = None + x._structure = None + x._annotation = None + + return x + + +def _are_types_castable(upstream: LiteralType, downstream: LiteralType) -> bool: + if upstream.collection_type is not None: + if upstream.collection_type is None: + return False + + return _are_types_castable(upstream.collection_type, downstream.collection_type) + + if upstream.map_value_type is not None: + if upstream.map_value_type is None: + return False + + return _are_types_castable(upstream.map_value_type, downstream.map_value_type) + + if upstream.structured_dataset_type is not None: + if downstream.structured_dataset_type is None: + return False + + usdt = upstream.structured_dataset_type + dsdt = downstream.structured_dataset_type + + if usdt.format != dsdt.format: + return False + + if usdt.external_schema_type != dsdt.external_schema_type: + return False + + if usdt.external_schema_bytes != dsdt.external_schema_bytes: + return False + + ucols = usdt.columns + dcols = dsdt.columns + + if len(ucols) != len(dcols): + return False + + for (u, d) in zip(ucols, dcols): + if u.name != d.name: + return False + + if not _are_types_castable(u.literal_type, d.literal_type): + return False + + return True + + if downstream.union_type is not None: + if upstream.union_type is not None: + # for each upstream variant, there must be a compatible type downstream + for v in upstream.union_type: + if not _are_types_castable(v, downstream): + return False + return True + + else: + # there must be a compatible downstream type + for v in downstream.union_type.variants: + if _are_types_castable(upstream, v): + return True + + if upstream.enum_type is not None: + # enums are castable to string + if downstream.simple == SimpleType.STRING: + return True + + if _type_essence(upstream) == _type_essence(downstream): + return True + + return False + + class UnionTransformer(TypeTransformer[T]): """ Transformer that handles a typing.Union[T1, T2, ...] @@ -880,6 +964,7 @@ def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], exp def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[T]) -> Optional[typing.Any]: union_tag = None + union_type = None if lv.scalar is not None and lv.scalar.union is not None: union_type = lv.scalar.union.stored_type if union_type.structure is not None: @@ -895,6 +980,10 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: if trans.name != union_tag: continue + expected_literal_type = TypeEngine.to_literal_type(v) + if not _are_types_castable(union_type, expected_literal_type): + continue + assert lv.scalar is not None # type checker assert lv.scalar.union is not None # type checker @@ -1014,8 +1103,11 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: # for empty generic we have to explicitly test for lv.scalar.generic is not None as empty dict # evaluates to false if lv and lv.scalar and lv.scalar.generic is not None: - return _json.loads(_json_format.MessageToJson(lv.scalar.generic)) - raise TypeError(f"Cannot convert from {lv} to {expected_python_type}") + try: + return _json.loads(_json_format.MessageToJson(lv.scalar.generic)) + except TypeError: + raise TypeTransformerFailedError(f"Cannot convert from {lv} to {expected_python_type}") + raise TypeTransformerFailedError(f"Cannot convert from {lv} to {expected_python_type}") def guess_python_type(self, literal_type: LiteralType) -> Type[T]: if literal_type.map_value_type: diff --git a/flytekit/types/file/file.py b/flytekit/types/file/file.py index b08bcd565c..8c5d4c6834 100644 --- a/flytekit/types/file/file.py +++ b/flytekit/types/file/file.py @@ -326,8 +326,10 @@ def to_literal( def to_python_value( self, ctx: FlyteContext, lv: Literal, expected_python_type: typing.Union[typing.Type[FlyteFile], os.PathLike] ) -> FlyteFile: - - uri = lv.scalar.blob.uri + try: + uri = lv.scalar.blob.uri + except AttributeError: + raise TypeTransformerFailedError(f"Cannot convert from {lv} to {expected_python_type}") # In this condition, we still return a FlyteFile instance, but it's a simple one that has no downloading tricks # Using is instead of issubclass because FlyteFile does actually subclass it if expected_python_type is os.PathLike: From 1d530a3f3a9df0950e0ac267d3717dc814ca867b Mon Sep 17 00:00:00 2001 From: Eduardo Apolinario Date: Thu, 3 Mar 2022 13:09:39 -0800 Subject: [PATCH 089/128] Point flyteidl to maxim's fork in CI and requirements files Signed-off-by: Eduardo Apolinario --- .github/workflows/pythonbuild.yml | 1 + dev-requirements.txt | 55 ++++++++----------- doc-requirements.txt | 32 +++++------ requirements-spark2.txt | 22 ++++---- requirements.in | 1 + requirements.txt | 22 ++++---- setup.py | 3 +- .../workflows/requirements.txt | 41 ++++++++++---- 8 files changed, 92 insertions(+), 85 deletions(-) diff --git a/.github/workflows/pythonbuild.yml b/.github/workflows/pythonbuild.yml index 56d26badd3..1d49e2010c 100644 --- a/.github/workflows/pythonbuild.yml +++ b/.github/workflows/pythonbuild.yml @@ -105,6 +105,7 @@ jobs: pip install -e . if [ -f dev-requirements.txt ]; then pip install -r dev-requirements.txt; fi pip install --no-deps -U https://github.com/flyteorg/flytekit/archive/${{ github.sha }}.zip#egg=flytekit + pip install --no-deps -U "git+https://github.com/maximsmol/flyteidl.git@maximsmol/union-types#egg=flyteidl" pip freeze - name: Test with coverage run: | diff --git a/dev-requirements.txt b/dev-requirements.txt index 6e17840325..43568eb535 100644 --- a/dev-requirements.txt +++ b/dev-requirements.txt @@ -32,6 +32,7 @@ certifi==2021.10.8 # requests cffi==1.15.0 # via + # -c requirements.txt # bcrypt # cryptography # pynacl @@ -41,7 +42,7 @@ chardet==4.0.0 # via # -c requirements.txt # binaryornot -charset-normalizer==2.0.11 +charset-normalizer==2.0.12 # via # -c requirements.txt # requests @@ -49,7 +50,7 @@ checksumdir==1.2.0 # via # -c requirements.txt # flytekit -click==8.0.3 +click==8.0.4 # via # -c requirements.txt # cookiecutter @@ -64,9 +65,9 @@ cookiecutter==1.7.3 # via # -c requirements.txt # flytekit -coverage[toml]==6.3.1 +coverage[toml]==6.3.2 # via -r dev-requirements.in -croniter==1.2.0 +croniter==1.3.4 # via # -c requirements.txt # flytekit @@ -93,7 +94,7 @@ diskcache==5.4.0 # flytekit distlib==0.3.4 # via virtualenv -distro==1.6.0 +distro==1.7.0 # via docker-compose docker[ssh]==5.0.3 # via docker-compose @@ -113,12 +114,8 @@ docstring-parser==0.13 # via # -c requirements.txt # flytekit -filelock==3.4.2 +filelock==3.6.0 # via virtualenv -flyteidl==0.22.0 - # via - # -c requirements.txt - # flytekit google-api-core[grpc]==2.5.0 # via # google-cloud-bigquery @@ -128,36 +125,36 @@ google-auth==2.6.0 # via # google-api-core # google-cloud-core -google-cloud-bigquery==2.32.0 +google-cloud-bigquery==2.34.1 # via -r dev-requirements.in -google-cloud-bigquery-storage==2.11.0 +google-cloud-bigquery-storage==2.12.0 # via -r dev-requirements.in google-cloud-core==2.2.2 # via google-cloud-bigquery google-crc32c==1.3.0 # via google-resumable-media -google-resumable-media==2.2.0 +google-resumable-media==2.3.1 # via google-cloud-bigquery -googleapis-common-protos==1.54.0 +googleapis-common-protos==1.55.0 # via # google-api-core # grpcio-status -grpcio==1.43.0 +grpcio==1.44.0 # via # -c requirements.txt # flytekit # google-api-core # google-cloud-bigquery # grpcio-status -grpcio-status==1.43.0 +grpcio-status==1.44.0 # via google-api-core -identify==2.4.8 +identify==2.4.11 # via pre-commit idna==3.3 # via # -c requirements.txt # requests -importlib-metadata==4.10.1 +importlib-metadata==4.11.2 # via # -c requirements.txt # keyring @@ -188,9 +185,7 @@ keyring==23.5.0 # via # -c requirements.txt # flytekit -libcst==0.4.1 - # via google-cloud-bigquery-storage -markupsafe==2.0.1 +markupsafe==2.1.0 # via # -c requirements.txt # jinja2 @@ -238,7 +233,7 @@ pandas==1.3.5 # flytekit paramiko==2.9.2 # via docker -platformdirs==2.4.1 +platformdirs==2.5.1 # via virtualenv pluggy==1.0.0 # via pytest @@ -248,14 +243,13 @@ poyo==0.5.0 # cookiecutter pre-commit==2.17.0 # via -r dev-requirements.in -proto-plus==1.20.0 +proto-plus==1.20.3 # via # google-cloud-bigquery # google-cloud-bigquery-storage protobuf==3.19.4 # via # -c requirements.txt - # flyteidl # flytekit # google-api-core # google-cloud-bigquery @@ -312,7 +306,7 @@ python-json-logger==2.0.2 # via # -c requirements.txt # flytekit -python-slugify==5.0.2 +python-slugify==6.1.1 # via # -c requirements.txt # cookiecutter @@ -329,9 +323,8 @@ pyyaml==5.4.1 # via # -c requirements.txt # docker-compose - # libcst # pre-commit -regex==2022.1.18 +regex==2022.3.2 # via # -c requirements.txt # docker-image-py @@ -389,29 +382,27 @@ toml==0.10.2 # via # pre-commit # pytest -tomli==2.0.0 +tomli==2.0.1 # via # coverage # mypy -typing-extensions==4.0.1 +typing-extensions==4.1.1 # via # -c requirements.txt # flytekit - # libcst # mypy # typing-inspect typing-inspect==0.7.1 # via # -c requirements.txt # dataclasses-json - # libcst urllib3==1.26.8 # via # -c requirements.txt # flytekit # requests # responses -virtualenv==20.13.1 +virtualenv==20.13.2 # via pre-commit websocket-client==0.59.0 # via diff --git a/doc-requirements.txt b/doc-requirements.txt index c320d836d5..caee82a1ed 100644 --- a/doc-requirements.txt +++ b/doc-requirements.txt @@ -10,7 +10,7 @@ alabaster==0.7.12 # via sphinx arrow==1.2.2 # via jinja2-time -astroid==2.9.3 +astroid==2.10.0 # via sphinx-autoapi babel==2.9.1 # via sphinx @@ -27,11 +27,11 @@ cffi==1.15.0 # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.11 +charset-normalizer==2.0.12 # via requests checksumdir==1.2.0 # via flytekit -click==8.0.3 +click==8.0.4 # via # cookiecutter # flytekit @@ -39,7 +39,7 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.2.0 +croniter==1.3.4 # via flytekit cryptography==36.0.1 # via @@ -63,11 +63,9 @@ docutils==0.17.1 # via # sphinx # sphinx-panels -flyteidl==0.22.0 - # via flytekit furo @ git+https://github.com/flyteorg/furo@main # via -r doc-requirements.in -grpcio==1.43.0 +grpcio==1.44.0 # via # -r doc-requirements.in # flytekit @@ -75,7 +73,7 @@ idna==3.3 # via requests imagesize==1.3.0 # via sphinx -importlib-metadata==4.10.1 +importlib-metadata==4.11.2 # via # keyring # sphinx @@ -95,9 +93,9 @@ keyring==23.5.0 # via flytekit lazy-object-proxy==1.7.1 # via astroid -lxml==4.7.1 +lxml==4.8.0 # via sphinx-material -markupsafe==2.0.1 +markupsafe==2.1.0 # via jinja2 marshmallow==3.14.1 # via @@ -118,14 +116,12 @@ numpy==1.22.2 # pyarrow packaging==21.3 # via sphinx -pandas==1.4.0 +pandas==1.4.1 # via flytekit poyo==0.5.0 # via cookiecutter protobuf==3.19.4 - # via - # flyteidl - # flytekit + # via flytekit py==1.11.0 # via retry pyarrow==6.0.1 @@ -146,7 +142,7 @@ python-dateutil==2.8.2 # pandas python-json-logger==2.0.2 # via flytekit -python-slugify[unidecode]==5.0.2 +python-slugify[unidecode]==6.1.1 # via # cookiecutter # sphinx-material @@ -159,7 +155,7 @@ pytz==2021.3 # pandas pyyaml==6.0 # via sphinx-autoapi -regex==2022.1.18 +regex==2022.3.2 # via docker-image-py requests==2.27.1 # via @@ -232,14 +228,14 @@ statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -typing-extensions==4.0.1 +typing-extensions==4.1.1 # via # astroid # flytekit # typing-inspect typing-inspect==0.7.1 # via dataclasses-json -unidecode==1.3.2 +unidecode==1.3.3 # via # python-slugify # sphinx-autoapi diff --git a/requirements-spark2.txt b/requirements-spark2.txt index 3e15827223..9ea67e42fc 100644 --- a/requirements-spark2.txt +++ b/requirements-spark2.txt @@ -4,6 +4,8 @@ # # make requirements-spark2.txt # +-e git+https://github.com/maximsmol/flyteidl.git@maximsmol/union-types#egg=flyteidl + # via -r requirements.in -e file:.#egg=flytekit # via # -r requirements-spark2.in @@ -22,11 +24,11 @@ cffi==1.15.0 # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.11 +charset-normalizer==2.0.12 # via requests checksumdir==1.2.0 # via flytekit -click==8.0.3 +click==8.0.4 # via # cookiecutter # flytekit @@ -34,7 +36,7 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.2.0 +croniter==1.3.4 # via flytekit cryptography==36.0.1 # via secretstorage @@ -50,13 +52,11 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.22.0 - # via flytekit -grpcio==1.43.0 +grpcio==1.44.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.10.1 +importlib-metadata==4.11.2 # via keyring jeepney==0.7.1 # via @@ -72,7 +72,7 @@ jsonschema==3.2.0 # via -r requirements.in keyring==23.5.0 # via flytekit -markupsafe==2.0.1 +markupsafe==2.1.0 # via jinja2 marshmallow==3.14.1 # via @@ -118,7 +118,7 @@ python-dateutil==2.8.2 # pandas python-json-logger==2.0.2 # via flytekit -python-slugify==5.0.2 +python-slugify==6.1.1 # via cookiecutter pytimeparse==1.1.8 # via flytekit @@ -128,7 +128,7 @@ pytz==2021.3 # pandas pyyaml==5.4.1 # via -r requirements.in -regex==2022.1.18 +regex==2022.3.2 # via docker-image-py requests==2.27.1 # via @@ -153,7 +153,7 @@ statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -typing-extensions==4.0.1 +typing-extensions==4.1.1 # via # flytekit # typing-inspect diff --git a/requirements.in b/requirements.in index 0cf01a6f13..545ef2c706 100644 --- a/requirements.in +++ b/requirements.in @@ -1,4 +1,5 @@ .[all] +-e git+https://github.com/maximsmol/flyteidl.git@maximsmol/union-types#egg=flyteidl -e file:.#egg=flytekit attrs<21 # We need to restrict constrain the versions of both jsonschema and pyyaml because of docker-compose (which is diff --git a/requirements.txt b/requirements.txt index d745fca2d9..8152af50b8 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,6 +4,8 @@ # # make requirements.txt # +-e git+https://github.com/maximsmol/flyteidl.git@maximsmol/union-types#egg=flyteidl + # via -r requirements.in -e file:.#egg=flytekit # via -r requirements.in arrow==1.2.2 @@ -20,11 +22,11 @@ cffi==1.15.0 # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.11 +charset-normalizer==2.0.12 # via requests checksumdir==1.2.0 # via flytekit -click==8.0.3 +click==8.0.4 # via # cookiecutter # flytekit @@ -32,7 +34,7 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.2.0 +croniter==1.3.4 # via flytekit cryptography==36.0.1 # via secretstorage @@ -48,13 +50,11 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.22.0 - # via flytekit -grpcio==1.43.0 +grpcio==1.44.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.10.1 +importlib-metadata==4.11.2 # via keyring jeepney==0.7.1 # via @@ -70,7 +70,7 @@ jsonschema==3.2.0 # via -r requirements.in keyring==23.5.0 # via flytekit -markupsafe==2.0.1 +markupsafe==2.1.0 # via jinja2 marshmallow==3.14.1 # via @@ -116,7 +116,7 @@ python-dateutil==2.8.2 # pandas python-json-logger==2.0.2 # via flytekit -python-slugify==5.0.2 +python-slugify==6.1.1 # via cookiecutter pytimeparse==1.1.8 # via flytekit @@ -126,7 +126,7 @@ pytz==2021.3 # pandas pyyaml==5.4.1 # via -r requirements.in -regex==2022.1.18 +regex==2022.3.2 # via docker-image-py requests==2.27.1 # via @@ -151,7 +151,7 @@ statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -typing-extensions==4.0.1 +typing-extensions==4.1.1 # via # flytekit # typing-inspect diff --git a/setup.py b/setup.py index 885f32b897..438fd9a1a6 100644 --- a/setup.py +++ b/setup.py @@ -35,7 +35,8 @@ ] }, install_requires=[ - "flyteidl>=0.22.0", + # TODO: put flyteidl back + # "flyteidl>=0.22.0", "wheel>=0.30.0,<1.0.0", "pandas>=1.0.0,<2.0.0", "pyarrow>=4.0.0,<7.0.0", diff --git a/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt b/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt index 93514898b2..fb365ed51d 100644 --- a/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt +++ b/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt @@ -10,13 +10,15 @@ binaryornot==0.4.4 # via cookiecutter certifi==2021.10.8 # via requests +cffi==1.15.0 + # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.11 +charset-normalizer==2.0.12 # via requests checksumdir==1.2.0 # via flytekit -click==7.1.2 +click==8.0.4 # via # cookiecutter # flytekit @@ -24,8 +26,10 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.2.0 +croniter==1.3.4 # via flytekit +cryptography==36.0.1 + # via secretstorage cycler==0.11.0 # via matplotlib dataclasses-json==0.5.6 @@ -40,18 +44,24 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.22.0 +flyteidl==0.23.0 # via flytekit -flytekit==0.30.0 +flytekit==0.30.3 # via -r tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.in fonttools==4.29.1 # via matplotlib -grpcio==1.43.0 +googleapis-common-protos==1.55.0 + # via flyteidl +grpcio==1.44.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.10.1 +importlib-metadata==4.11.2 # via keyring +jeepney==0.7.1 + # via + # keyring + # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -64,7 +74,7 @@ keyring==23.5.0 # via flytekit kiwisolver==1.3.2 # via matplotlib -markupsafe==2.0.1 +markupsafe==2.1.0 # via jinja2 marshmallow==3.14.1 # via @@ -91,7 +101,7 @@ opencv-python==4.5.5.62 # via -r tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.in packaging==21.3 # via matplotlib -pandas==1.4.0 +pandas==1.4.1 # via flytekit pillow==9.0.1 # via matplotlib @@ -101,6 +111,10 @@ protobuf==3.19.4 # via # flyteidl # flytekit + # googleapis-common-protos + # protoc-gen-swagger +protoc-gen-swagger==0.1.0 + # via flyteidl py==1.11.0 # via retry pyarrow==6.0.1 @@ -113,13 +127,14 @@ pyparsing==3.0.7 # packaging python-dateutil==2.8.2 # via + # arrow # croniter # flytekit # matplotlib # pandas python-json-logger==2.0.2 # via flytekit -python-slugify==5.0.2 +python-slugify==6.1.1 # via cookiecutter pytimeparse==1.1.8 # via flytekit @@ -127,7 +142,7 @@ pytz==2021.3 # via # flytekit # pandas -regex==2022.1.18 +regex==2022.3.2 # via docker-image-py requests==2.27.1 # via @@ -138,6 +153,8 @@ responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit +secretstorage==3.3.1 + # via keyring six==1.16.0 # via # cookiecutter @@ -149,7 +166,7 @@ statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -typing-extensions==4.0.1 +typing-extensions==4.1.1 # via # flytekit # typing-inspect From 35de6973a1bdaa10919b695a536dd15644c69f73 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Fri, 18 Feb 2022 03:32:00 +0800 Subject: [PATCH 090/128] [Core feature] Add Raw AWS Batch Task (#782) * Init plugin Signed-off-by: Kevin Su * Fixed lint Signed-off-by: Kevin Su * address comment Signed-off-by: Kevin Su * Fixed typo Signed-off-by: Kevin Su * Updated AWS config * Fixed lint Signed-off-by: Kevin Su * Added comment Signed-off-by: Kevin Su * Update config Signed-off-by: Kevin Su * Fixed tests Signed-off-by: Kevin Su * Fixed tests Signed-off-by: Kevin Su * Fixed tests Signed-off-by: Kevin Su * use pyflyte execute Signed-off-by: Kevin Su * Fixed tests Signed-off-by: Kevin Su * Fixed tests Signed-off-by: Kevin Su * Added comment Signed-off-by: Kevin Su Signed-off-by: maximsmol --- .github/workflows/pythonbuild.yml | 1 + flytekit/bin/entrypoint.py | 25 ++- flytekit/core/map_task.py | 2 +- plugins/flytekit-aws-batch/README.md | 9 ++ .../flytekitplugins/awsbatch/__init__.py | 1 + .../flytekitplugins/awsbatch/task.py | 80 ++++++++++ plugins/flytekit-aws-batch/requirements.in | 2 + plugins/flytekit-aws-batch/requirements.txt | 148 ++++++++++++++++++ plugins/flytekit-aws-batch/setup.py | 36 +++++ plugins/flytekit-aws-batch/tests/__init__.py | 0 .../tests/test_aws_batch.py | 49 ++++++ 11 files changed, 351 insertions(+), 2 deletions(-) create mode 100644 plugins/flytekit-aws-batch/README.md create mode 100644 plugins/flytekit-aws-batch/flytekitplugins/awsbatch/__init__.py create mode 100644 plugins/flytekit-aws-batch/flytekitplugins/awsbatch/task.py create mode 100644 plugins/flytekit-aws-batch/requirements.in create mode 100644 plugins/flytekit-aws-batch/requirements.txt create mode 100644 plugins/flytekit-aws-batch/setup.py create mode 100644 plugins/flytekit-aws-batch/tests/__init__.py create mode 100644 plugins/flytekit-aws-batch/tests/test_aws_batch.py diff --git a/.github/workflows/pythonbuild.yml b/.github/workflows/pythonbuild.yml index 1d49e2010c..f1cfa4977f 100644 --- a/.github/workflows/pythonbuild.yml +++ b/.github/workflows/pythonbuild.yml @@ -64,6 +64,7 @@ jobs: plugin-names: # Please maintain an alphabetical order in the following list - flytekit-aws-athena + - flytekit-aws-batch - flytekit-aws-sagemaker - flytekit-bigquery - flytekit-data-fsspec diff --git a/flytekit/bin/entrypoint.py b/flytekit/bin/entrypoint.py index 4fc9f69d3e..9d07d38a63 100644 --- a/flytekit/bin/entrypoint.py +++ b/flytekit/bin/entrypoint.py @@ -58,7 +58,9 @@ def _compute_array_job_index(): offset = 0 if _os.environ.get("BATCH_JOB_ARRAY_INDEX_OFFSET"): offset = int(_os.environ.get("BATCH_JOB_ARRAY_INDEX_OFFSET")) - return offset + int(_os.environ.get(_os.environ.get("BATCH_JOB_ARRAY_INDEX_VAR_NAME"))) + if _os.environ.get("BATCH_JOB_ARRAY_INDEX_VAR_NAME"): + return offset + int(_os.environ.get(_os.environ.get("BATCH_JOB_ARRAY_INDEX_VAR_NAME"))) + return offset def _dispatch_execute( @@ -345,6 +347,27 @@ def _execute_map_task( dynamic_addl_distro: Optional[str] = None, dynamic_dest_dir: Optional[str] = None, ): + """ + This function should be called by map task and aws-batch task + resolver should be something like: + flytekit.core.python_auto_container.default_task_resolver + resolver args should be something like + task_module app.workflows task_name task_1 + have dashes seems to mess up click, like --task_module seems to interfere + + :param inputs: Where to read inputs + :param output_prefix: Where to write primitive outputs + :param raw_output_data_prefix: Where to write offloaded data (files, directories, dataframes). + :param test: Dry run + :param resolver: The task resolver to use. This needs to be loadable directly from importlib (and thus cannot be + nested). + :param resolver_args: Args that will be passed to the aforementioned resolver's load_task function + :param dynamic_addl_distro: In the case of parent tasks executed using the 'fast' mode this captures where the + compressed code archive has been uploaded. + :param dynamic_dest_dir: In the case of parent tasks executed using the 'fast' mode this captures where compressed + code archives should be installed in the flyte task container. + :return: + """ if len(resolver_args) < 1: raise Exception(f"Resolver args cannot be <1, got {resolver_args}") diff --git a/flytekit/core/map_task.py b/flytekit/core/map_task.py index 4ec9f64a2d..731be53ba8 100644 --- a/flytekit/core/map_task.py +++ b/flytekit/core/map_task.py @@ -52,7 +52,7 @@ def __init__( collection_interface = transform_interface_to_list_interface(python_function_task.python_interface) instance = next(self._ids) - name = f"{python_function_task._task_function.__module__}.mapper_{python_function_task._task_function.__name__}_{instance}" + name = f"{python_function_task.task_function.__module__}.mapper_{python_function_task.task_function.__name__}_{instance}" self._run_task = python_function_task self._max_concurrency = concurrency diff --git a/plugins/flytekit-aws-batch/README.md b/plugins/flytekit-aws-batch/README.md new file mode 100644 index 0000000000..b56663237b --- /dev/null +++ b/plugins/flytekit-aws-batch/README.md @@ -0,0 +1,9 @@ +# Flytekit AWS Batch Plugin + +Flyte backend can be connected with AWS batch. Once enabled, it allows you to run flyte task on AWS batch service + +To install the plugin, run the following command: + +```bash +pip install flytekitplugins-awsbatch +``` diff --git a/plugins/flytekit-aws-batch/flytekitplugins/awsbatch/__init__.py b/plugins/flytekit-aws-batch/flytekitplugins/awsbatch/__init__.py new file mode 100644 index 0000000000..244716ebe7 --- /dev/null +++ b/plugins/flytekit-aws-batch/flytekitplugins/awsbatch/__init__.py @@ -0,0 +1 @@ +from .task import AWSBatchConfig diff --git a/plugins/flytekit-aws-batch/flytekitplugins/awsbatch/task.py b/plugins/flytekit-aws-batch/flytekitplugins/awsbatch/task.py new file mode 100644 index 0000000000..c9b30b7af1 --- /dev/null +++ b/plugins/flytekit-aws-batch/flytekitplugins/awsbatch/task.py @@ -0,0 +1,80 @@ +from dataclasses import dataclass +from typing import Any, Callable, Dict, List, Optional + +from dataclasses_json import dataclass_json +from google.protobuf import json_format +from google.protobuf.struct_pb2 import Struct + +from flytekit import PythonFunctionTask +from flytekit.extend import SerializationSettings, TaskPlugins + + +@dataclass_json +@dataclass +class AWSBatchConfig(object): + """ + Use this to configure SubmitJobInput for a AWS batch job. Task's marked with this will automatically execute + natively onto AWS batch service. + Refer to AWS SubmitJobInput for more detail: https://docs.aws.amazon.com/sdk-for-go/api/service/batch/#SubmitJobInput + """ + + parameters: Optional[Dict[str, str]] = None + schedulingPriority: Optional[int] = None + platformCapabilities: str = "EC2" + propagateTags: Optional[bool] = None + tags: Optional[Dict[str, str]] = None + + def to_dict(self): + s = Struct() + s.update(self.to_dict()) + return json_format.MessageToDict(s) + + +class AWSBatchFunctionTask(PythonFunctionTask): + """ + Actual Plugin that transforms the local python code for execution within AWS batch job + """ + + _AWS_BATCH_TASK_TYPE = "aws-batch" + + def __init__(self, task_config: AWSBatchConfig, task_function: Callable, **kwargs): + if task_config is None: + task_config = AWSBatchConfig() + super(AWSBatchFunctionTask, self).__init__( + task_config=task_config, task_type=self._AWS_BATCH_TASK_TYPE, task_function=task_function, **kwargs + ) + self._task_config = task_config + + def get_custom(self, settings: SerializationSettings) -> Dict[str, Any]: + # task_config will be used to create SubmitJobInput in propeller except platformCapabilities. + return self._task_config.to_dict() + + def get_config(self, settings: SerializationSettings) -> Dict[str, str]: + # Parameters in taskTemplate config will be used to create aws job definition. + # More detail about job definition: https://docs.aws.amazon.com/batch/latest/userguide/job_definition_parameters.html + return {"platformCapabilities": self._task_config.platformCapabilities} + + def get_command(self, settings: SerializationSettings) -> List[str]: + container_args = [ + "pyflyte-execute", + "--inputs", + "{{.input}}", + "--output-prefix", + # As of FlytePropeller v0.16.28, aws array batch plugin support to run single job. + # This task will call aws batch plugin to execute the task on aws batch service. + # For single job, FlytePropeller will always read the output from this directory (outputPrefix/0) + # More detail, see https://github.com/flyteorg/flyteplugins/blob/0dd93c23ed2edeca65d58e89b0edb613f88120e0/go/tasks/plugins/array/catalog.go#L501. + "{{.outputPrefix}}/0", + "--raw-output-data-prefix", + "{{.rawOutputDataPrefix}}", + "--resolver", + self.task_resolver.location, + "--", + *self.task_resolver.loader_args(settings, self), + ] + + return container_args + + +# Inject the AWS batch plugin into flytekits dynamic plugin loading system +TaskPlugins.register_pythontask_plugin(AWSBatchConfig, AWSBatchFunctionTask) diff --git a/plugins/flytekit-aws-batch/requirements.in b/plugins/flytekit-aws-batch/requirements.in new file mode 100644 index 0000000000..4c2c1fae97 --- /dev/null +++ b/plugins/flytekit-aws-batch/requirements.in @@ -0,0 +1,2 @@ +. +-e file:.#egg=flytekitplugins-awsbatch diff --git a/plugins/flytekit-aws-batch/requirements.txt b/plugins/flytekit-aws-batch/requirements.txt new file mode 100644 index 0000000000..9181490234 --- /dev/null +++ b/plugins/flytekit-aws-batch/requirements.txt @@ -0,0 +1,148 @@ +# +# This file is autogenerated by pip-compile with python 3.9 +# To update, run: +# +# pip-compile requirements.in +# +-e file:.#egg=flytekitplugins-awsbatch + # via -r requirements.in +arrow==1.2.1 + # via jinja2-time +binaryornot==0.4.4 + # via cookiecutter +certifi==2021.10.8 + # via requests +chardet==4.0.0 + # via binaryornot +charset-normalizer==2.0.10 + # via requests +checksumdir==1.2.0 + # via flytekit +click==7.1.2 + # via + # cookiecutter + # flytekit +cloudpickle==2.0.0 + # via flytekit +cookiecutter==1.7.3 + # via flytekit +croniter==1.2.0 + # via flytekit +dataclasses-json==0.5.6 + # via flytekit +decorator==5.1.1 + # via retry +deprecated==1.2.13 + # via flytekit +diskcache==5.4.0 + # via flytekit +docker-image-py==0.1.12 + # via flytekit +docstring-parser==0.13 + # via flytekit +flyteidl==0.21.23 + # via flytekit +flytekit==0.26.0 + # via flytekitplugins-awsbatch +grpcio==1.43.0 + # via flytekit +idna==3.3 + # via requests +importlib-metadata==4.10.1 + # via keyring +jinja2==3.0.3 + # via + # cookiecutter + # jinja2-time +jinja2-time==0.2.0 + # via cookiecutter +keyring==23.5.0 + # via flytekit +markupsafe==2.0.1 + # via jinja2 +marshmallow==3.14.1 + # via + # dataclasses-json + # marshmallow-enum + # marshmallow-jsonschema +marshmallow-enum==1.5.1 + # via dataclasses-json +marshmallow-jsonschema==0.13.0 + # via flytekit +mypy-extensions==0.4.3 + # via typing-inspect +natsort==8.0.2 + # via flytekit +numpy==1.22.1 + # via + # pandas + # pyarrow +pandas==1.3.5 + # via flytekit +poyo==0.5.0 + # via cookiecutter +protobuf==3.19.3 + # via + # flyteidl + # flytekit +py==1.11.0 + # via retry +pyarrow==6.0.1 + # via flytekit +python-dateutil==2.8.1 + # via + # arrow + # croniter + # flytekit + # pandas +python-json-logger==2.0.2 + # via flytekit +python-slugify==5.0.2 + # via cookiecutter +pytimeparse==1.1.8 + # via flytekit +pytz==2021.3 + # via + # flytekit + # pandas +regex==2022.1.18 + # via docker-image-py +requests==2.27.1 + # via + # cookiecutter + # flytekit + # responses +responses==0.17.0 + # via flytekit +retry==0.9.2 + # via flytekit +six==1.16.0 + # via + # cookiecutter + # flytekit + # grpcio + # python-dateutil + # responses +sortedcontainers==2.4.0 + # via flytekit +statsd==3.3.0 + # via flytekit +text-unidecode==1.3 + # via python-slugify +typing-extensions==4.0.1 + # via typing-inspect +typing-inspect==0.7.1 + # via dataclasses-json +urllib3==1.26.8 + # via + # flytekit + # requests + # responses +wheel==0.37.1 + # via flytekit +wrapt==1.13.3 + # via + # deprecated + # flytekit +zipp==3.7.0 + # via importlib-metadata diff --git a/plugins/flytekit-aws-batch/setup.py b/plugins/flytekit-aws-batch/setup.py new file mode 100644 index 0000000000..43613fe244 --- /dev/null +++ b/plugins/flytekit-aws-batch/setup.py @@ -0,0 +1,36 @@ +from setuptools import setup + +PLUGIN_NAME = "awsbatch" + +microlib_name = f"flytekitplugins-{PLUGIN_NAME}" + +plugin_requires = ["flytekit>=0.19.0,<1.0.0"] + +__version__ = "0.0.0+develop" + +setup( + name=microlib_name, + version=__version__, + author="flyteorg", + author_email="admin@flyte.org", + description="This package holds the AWS Batch plugins for flytekit", + namespace_packages=["flytekitplugins"], + packages=[f"flytekitplugins.{PLUGIN_NAME}"], + install_requires=plugin_requires, + license="apache2", + python_requires=">=3.7", + classifiers=[ + "Intended Audience :: Science/Research", + "Intended Audience :: Developers", + "License :: OSI Approved :: Apache Software License", + "Programming Language :: Python :: 3.7", + "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Topic :: Scientific/Engineering", + "Topic :: Scientific/Engineering :: Artificial Intelligence", + "Topic :: Software Development", + "Topic :: Software Development :: Libraries", + "Topic :: Software Development :: Libraries :: Python Modules", + ], +) diff --git a/plugins/flytekit-aws-batch/tests/__init__.py b/plugins/flytekit-aws-batch/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/plugins/flytekit-aws-batch/tests/test_aws_batch.py b/plugins/flytekit-aws-batch/tests/test_aws_batch.py new file mode 100644 index 0000000000..5cd1e5e5c2 --- /dev/null +++ b/plugins/flytekit-aws-batch/tests/test_aws_batch.py @@ -0,0 +1,49 @@ +from flytekitplugins.awsbatch import AWSBatchConfig + +from flytekit import PythonFunctionTask, task +from flytekit.extend import Image, ImageConfig, SerializationSettings + +config = AWSBatchConfig( + parameters={"codec": "mp4"}, + platformCapabilities="EC2", + propagateTags=True, + tags={"hello": "world"}, +) + + +def test_aws_batch_task(): + @task(task_config=config) + def t1(a: int) -> str: + inc = a + 2 + return str(inc) + + assert t1.task_config is not None + assert t1.task_config == config + assert t1.task_type == "aws-batch" + assert isinstance(t1, PythonFunctionTask) + + default_img = Image(name="default", fqn="test", tag="tag") + settings = SerializationSettings( + project="project", + domain="domain", + version="version", + env={"FOO": "baz"}, + image_config=ImageConfig(default_image=default_img, images=[default_img]), + ) + assert t1.get_custom(settings) == config.to_dict() + assert t1.get_command(settings) == [ + "pyflyte-execute", + "--inputs", + "{{.input}}", + "--output-prefix", + "{{.outputPrefix}}/0", + "--raw-output-data-prefix", + "{{.rawOutputDataPrefix}}", + "--resolver", + "flytekit.core.python_auto_container.default_task_resolver", + "--", + "task-module", + "tests.test_aws_batch", + "task-name", + "t1", + ] From d5823d267909da4e8db1a9de12f8561d34ff332f Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Fri, 18 Feb 2022 04:55:44 +0800 Subject: [PATCH 091/128] Add structured dataset encoder/decoder in fsspec plugin (#849) Signed-off-by: Kevin Su Signed-off-by: maximsmol --- .../flytekit-data-fsspec/dev-requirements.in | 1 + .../flytekit-data-fsspec/dev-requirements.txt | 52 +++++++++++++ .../flytekitplugins/fsspec/__init__.py | 19 +++++ .../flytekitplugins/fsspec/arrow.py | 69 +++++++++++++++++ .../flytekitplugins/fsspec/pandas.py | 74 +++++++++++++++++++ .../flytekitplugins/fsspec/persist.py | 43 ++++++----- plugins/flytekit-data-fsspec/requirements.txt | 10 ++- plugins/flytekit-data-fsspec/setup.py | 6 +- .../tests/test_basic_dfs.py | 44 +++++++++++ .../tests/test_persist.py | 22 ++++-- 10 files changed, 309 insertions(+), 31 deletions(-) create mode 100644 plugins/flytekit-data-fsspec/dev-requirements.in create mode 100644 plugins/flytekit-data-fsspec/dev-requirements.txt create mode 100644 plugins/flytekit-data-fsspec/flytekitplugins/fsspec/arrow.py create mode 100644 plugins/flytekit-data-fsspec/flytekitplugins/fsspec/pandas.py create mode 100644 plugins/flytekit-data-fsspec/tests/test_basic_dfs.py diff --git a/plugins/flytekit-data-fsspec/dev-requirements.in b/plugins/flytekit-data-fsspec/dev-requirements.in new file mode 100644 index 0000000000..a51391a948 --- /dev/null +++ b/plugins/flytekit-data-fsspec/dev-requirements.in @@ -0,0 +1 @@ +s3fs diff --git a/plugins/flytekit-data-fsspec/dev-requirements.txt b/plugins/flytekit-data-fsspec/dev-requirements.txt new file mode 100644 index 0000000000..b91ba07e85 --- /dev/null +++ b/plugins/flytekit-data-fsspec/dev-requirements.txt @@ -0,0 +1,52 @@ +# +# This file is autogenerated by pip-compile with python 3.9 +# To update, run: +# +# pip-compile dev-requirements.in +# +aiobotocore==2.1.1 + # via s3fs +aiohttp==3.8.1 + # via + # aiobotocore + # s3fs +aioitertools==0.9.0 + # via aiobotocore +aiosignal==1.2.0 + # via aiohttp +async-timeout==4.0.2 + # via aiohttp +attrs==21.4.0 + # via aiohttp +botocore==1.23.24 + # via aiobotocore +charset-normalizer==2.0.12 + # via aiohttp +frozenlist==1.3.0 + # via + # aiohttp + # aiosignal +fsspec==2022.01.0 + # via s3fs +idna==3.3 + # via yarl +jmespath==0.10.0 + # via botocore +multidict==6.0.2 + # via + # aiohttp + # yarl +python-dateutil==2.8.2 + # via botocore +s3fs==2022.1.0 + # via -r dev-requirements.in +six==1.16.0 + # via python-dateutil +typing-extensions==4.1.1 + # via aioitertools +urllib3==1.26.8 + # via botocore +wrapt==1.13.3 + # via aiobotocore +yarl==1.7.2 + # via aiohttp diff --git a/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/__init__.py b/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/__init__.py index c82c82792e..85a65d17b8 100644 --- a/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/__init__.py +++ b/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/__init__.py @@ -1 +1,20 @@ +import importlib + +from flytekit import USE_STRUCTURED_DATASET, StructuredDatasetTransformerEngine, logger +from flytekit.types.structured.structured_dataset import S3 + from .persist import FSSpecPersistence + +if USE_STRUCTURED_DATASET.get(): + from .arrow import ArrowToParquetEncodingHandler, ParquetToArrowDecodingHandler + from .pandas import PandasToParquetEncodingHandler, ParquetToPandasDecodingHandler + + def _register(protocol: str): + logger.info(f"Registering fsspec {protocol} implementations and overriding default structured encoder/decoder.") + StructuredDatasetTransformerEngine.register(PandasToParquetEncodingHandler(protocol), True, True) + StructuredDatasetTransformerEngine.register(ParquetToPandasDecodingHandler(protocol), True, True) + StructuredDatasetTransformerEngine.register(ArrowToParquetEncodingHandler(protocol), True, True) + StructuredDatasetTransformerEngine.register(ParquetToArrowDecodingHandler(protocol), True, True) + + if importlib.util.find_spec("s3fs"): + _register(S3) diff --git a/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/arrow.py b/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/arrow.py new file mode 100644 index 0000000000..c17e2fc8bd --- /dev/null +++ b/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/arrow.py @@ -0,0 +1,69 @@ +import os +import typing +from pathlib import Path + +import pyarrow as pa +import pyarrow.parquet as pq +from botocore.exceptions import NoCredentialsError +from flytekitplugins.fsspec.persist import FSSpecPersistence +from fsspec.core import split_protocol, strip_protocol + +from flytekit import FlyteContext, logger +from flytekit.models import literals +from flytekit.models.literals import StructuredDatasetMetadata +from flytekit.models.types import StructuredDatasetType +from flytekit.types.structured.structured_dataset import ( + PARQUET, + StructuredDataset, + StructuredDatasetDecoder, + StructuredDatasetEncoder, +) + + +class ArrowToParquetEncodingHandler(StructuredDatasetEncoder): + def __init__(self, protocol: str): + super().__init__(pa.Table, protocol, PARQUET) + + def encode( + self, + ctx: FlyteContext, + structured_dataset: StructuredDataset, + structured_dataset_type: StructuredDatasetType, + ) -> literals.StructuredDataset: + uri = typing.cast(str, structured_dataset.uri) or ctx.file_access.get_random_remote_directory() + if not ctx.file_access.is_remote(uri): + Path(uri).mkdir(parents=True, exist_ok=True) + path = os.path.join(uri, f"{0:05}") + filesystem = FSSpecPersistence.get_filesystem(path) + pq.write_table(structured_dataset.dataframe, strip_protocol(path), filesystem=filesystem) + return literals.StructuredDataset(uri=uri, metadata=StructuredDatasetMetadata(structured_dataset_type)) + + +class ParquetToArrowDecodingHandler(StructuredDatasetDecoder): + def __init__(self, protocol: str): + super().__init__(pa.Table, protocol, PARQUET) + + def decode( + self, + ctx: FlyteContext, + flyte_value: literals.StructuredDataset, + ) -> pa.Table: + uri = flyte_value.uri + if not ctx.file_access.is_remote(uri): + Path(uri).parent.mkdir(parents=True, exist_ok=True) + _, path = split_protocol(uri) + + columns = None + if flyte_value.metadata.structured_dataset_type.columns: + columns = [] + for c in flyte_value.metadata.structured_dataset_type.columns: + columns.append(c.name) + try: + fs = FSSpecPersistence.get_filesystem(uri) + return pq.read_table(path, filesystem=fs, columns=columns) + except NoCredentialsError as e: + logger.debug("S3 source detected, attempting anonymous S3 access") + fs = FSSpecPersistence.get_anonymous_filesystem(uri) + if fs is not None: + return pq.read_table(path, filesystem=fs, columns=columns) + raise e diff --git a/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/pandas.py b/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/pandas.py new file mode 100644 index 0000000000..52bcc4522a --- /dev/null +++ b/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/pandas.py @@ -0,0 +1,74 @@ +import os +import typing +from pathlib import Path + +import pandas as pd +from botocore.exceptions import NoCredentialsError +from flytekitplugins.fsspec.persist import FSSpecPersistence, s3_setup_args + +from flytekit import FlyteContext, logger +from flytekit.models import literals +from flytekit.models.literals import StructuredDatasetMetadata +from flytekit.models.types import StructuredDatasetType +from flytekit.types.structured.structured_dataset import ( + PARQUET, + S3, + StructuredDataset, + StructuredDatasetDecoder, + StructuredDatasetEncoder, +) + + +def get_storage_options(uri: str) -> typing.Optional[typing.Dict]: + protocol = FSSpecPersistence.get_protocol(uri) + if protocol == S3: + kwargs = s3_setup_args() + if kwargs: + return kwargs + return None + + +class PandasToParquetEncodingHandler(StructuredDatasetEncoder): + def __init__(self, protocol: str): + super().__init__(pd.DataFrame, protocol, PARQUET) + + def encode( + self, + ctx: FlyteContext, + structured_dataset: StructuredDataset, + structured_dataset_type: StructuredDatasetType, + ) -> literals.StructuredDataset: + uri = typing.cast(str, structured_dataset.uri) or ctx.file_access.get_random_remote_directory() + if not ctx.file_access.is_remote(uri): + Path(uri).mkdir(parents=True, exist_ok=True) + path = os.path.join(uri, f"{0:05}") + df = typing.cast(pd.DataFrame, structured_dataset.dataframe) + df.to_parquet( + path, coerce_timestamps="us", allow_truncated_timestamps=False, storage_options=get_storage_options(path) + ) + structured_dataset_type.format = PARQUET + return literals.StructuredDataset(uri=uri, metadata=StructuredDatasetMetadata(structured_dataset_type)) + + +class ParquetToPandasDecodingHandler(StructuredDatasetDecoder): + def __init__(self, protocol: str): + super().__init__(pd.DataFrame, protocol, PARQUET) + + def decode( + self, + ctx: FlyteContext, + flyte_value: literals.StructuredDataset, + ) -> pd.DataFrame: + uri = flyte_value.uri + columns = None + kwargs = get_storage_options(uri) + if flyte_value.metadata.structured_dataset_type.columns: + columns = [] + for c in flyte_value.metadata.structured_dataset_type.columns: + columns.append(c.name) + try: + return pd.read_parquet(uri, columns=columns, storage_options=kwargs) + except NoCredentialsError: + logger.debug("S3 source detected, attempting anonymous S3 access") + kwargs["anon"] = True + return pd.read_parquet(uri, columns=columns, storage_options=kwargs) diff --git a/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/persist.py b/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/persist.py index 68bc92b493..8583f47a91 100644 --- a/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/persist.py +++ b/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/persist.py @@ -35,29 +35,38 @@ class FSSpecPersistence(DataPersistence): def __init__(self, default_prefix=None): super(FSSpecPersistence, self).__init__(name="fsspec-persistence", default_prefix=default_prefix) - self.default_protocol = self._get_protocol(default_prefix) + self.default_protocol = self.get_protocol(default_prefix) @staticmethod - def _get_protocol(path: typing.Optional[str] = None): + def get_protocol(path: typing.Optional[str] = None): if path: protocol, _ = split_protocol(path) if protocol is None and path.startswith("/"): - print("Setting protocol to file") + logger.info("Setting protocol to file") protocol = "file" else: protocol = "file" return protocol @staticmethod - def _get_filesystem(path: str) -> fsspec.AbstractFileSystem: - protocol = FSSpecPersistence._get_protocol(path) + def get_filesystem(path: str) -> fsspec.AbstractFileSystem: + protocol = FSSpecPersistence.get_protocol(path) kwargs = {} if protocol == "file": kwargs = {"auto_mkdir": True} - if protocol == "s3": + elif protocol == "s3": kwargs = s3_setup_args() return fsspec.filesystem(protocol, **kwargs) # type: ignore + @staticmethod + def get_anonymous_filesystem(path: str) -> typing.Optional[fsspec.AbstractFileSystem]: + protocol = FSSpecPersistence.get_protocol(path) + if protocol == "s3": + kwargs = s3_setup_args() + anonymous_fs = fsspec.filesystem(protocol, anon=True, **kwargs) # type: ignore + return anonymous_fs + return None + @staticmethod def recursive_paths(f: str, t: str) -> typing.Tuple[str, str]: if not f.endswith("*"): @@ -68,36 +77,32 @@ def recursive_paths(f: str, t: str) -> typing.Tuple[str, str]: def exists(self, path: str) -> bool: try: - fs = self._get_filesystem(path) + fs = self.get_filesystem(path) return fs.exists(path) except OSError as oe: logger.debug(f"Error in exists checking {path} {oe}") - protocol = FSSpecPersistence._get_protocol(path) - if protocol == "s3": + fs = self.get_anonymous_filesystem(path) + if fs is not None: logger.debug("S3 source detected, attempting anonymous S3 exists check") - kwargs = s3_setup_args() - anonymous_fs = fsspec.filesystem(protocol, anon=True, **kwargs) # type: ignore - return anonymous_fs.exists(path) + return fs.exists(path) raise oe def get(self, from_path: str, to_path: str, recursive: bool = False): - fs = self._get_filesystem(from_path) + fs = self.get_filesystem(from_path) if recursive: from_path, to_path = self.recursive_paths(from_path, to_path) try: return fs.get(from_path, to_path, recursive=recursive) except OSError as oe: logger.debug(f"Error in getting {from_path} to {to_path} rec {recursive} {oe}") - protocol = FSSpecPersistence._get_protocol(from_path) - if protocol == "s3": + fs = self.get_anonymous_filesystem(from_path) + if fs is not None: logger.debug("S3 source detected, attempting anonymous S3 access") - kwargs = s3_setup_args() - anonymous_fs = fsspec.filesystem(protocol, anon=True, **kwargs) # type: ignore - return anonymous_fs.get(from_path, to_path, recursive=recursive) + return fs.get(from_path, to_path, recursive=recursive) raise oe def put(self, from_path: str, to_path: str, recursive: bool = False): - fs = self._get_filesystem(to_path) + fs = self.get_filesystem(to_path) if recursive: from_path, to_path = self.recursive_paths(from_path, to_path) # BEGIN HACK! diff --git a/plugins/flytekit-data-fsspec/requirements.txt b/plugins/flytekit-data-fsspec/requirements.txt index c644ac98e7..4dc89fa500 100644 --- a/plugins/flytekit-data-fsspec/requirements.txt +++ b/plugins/flytekit-data-fsspec/requirements.txt @@ -10,6 +10,8 @@ arrow==1.2.2 # via jinja2-time binaryornot==0.4.4 # via cookiecutter +botocore==1.23.24 + # via flytekitplugins-data-fsspec certifi==2021.10.8 # via requests chardet==4.0.0 @@ -28,8 +30,6 @@ cookiecutter==1.7.3 # via flytekit croniter==1.2.0 # via flytekit -cryptography==36.0.1 - # via secretstorage dataclasses-json==0.5.6 # via flytekit decorator==5.1.1 @@ -60,6 +60,8 @@ jinja2==3.0.3 # jinja2-time jinja2-time==0.2.0 # via cookiecutter +jmespath==0.10.0 + # via botocore keyring==23.5.0 # via flytekit markupsafe==2.0.1 @@ -93,11 +95,10 @@ py==1.11.0 # via retry pyarrow==6.0.1 # via flytekit -pycparser==2.21 - # via cffi python-dateutil==2.8.2 # via # arrow + # botocore # croniter # flytekit # pandas @@ -141,6 +142,7 @@ typing-inspect==0.7.1 # via dataclasses-json urllib3==1.26.8 # via + # botocore # flytekit # requests # responses diff --git a/plugins/flytekit-data-fsspec/setup.py b/plugins/flytekit-data-fsspec/setup.py index 3830d9d168..a0986a2e82 100644 --- a/plugins/flytekit-data-fsspec/setup.py +++ b/plugins/flytekit-data-fsspec/setup.py @@ -4,7 +4,7 @@ microlib_name = f"flytekitplugins-data-{PLUGIN_NAME}" -plugin_requires = ["flytekit>=0.21.3,<1.0.0", "fsspec>=2021.7.0"] +plugin_requires = ["flytekit>=0.21.3,<1.0.0", "fsspec>=2021.7.0", "botocore>=1.7.48"] __version__ = "0.0.0+develop" @@ -20,6 +20,10 @@ namespace_packages=["flytekitplugins"], packages=[f"flytekitplugins.{PLUGIN_NAME}"], install_requires=plugin_requires, + extras_require={ + # https://github.com/fsspec/filesystem_spec/blob/master/setup.py#L36 + "aws": ["s3fs>=2021.7.0"], + }, license="apache2", python_requires=">=3.7", classifiers=[ diff --git a/plugins/flytekit-data-fsspec/tests/test_basic_dfs.py b/plugins/flytekit-data-fsspec/tests/test_basic_dfs.py new file mode 100644 index 0000000000..b2de2cb3e2 --- /dev/null +++ b/plugins/flytekit-data-fsspec/tests/test_basic_dfs.py @@ -0,0 +1,44 @@ +import pandas as pd +import pyarrow as pa +from flytekitplugins.fsspec.pandas import get_storage_options + +from flytekit import kwtypes, task +from flytekit.configuration import aws + +try: + from typing import Annotated +except ImportError: + from typing_extensions import Annotated + + +def test_get_storage_options(): + endpoint = "https://s3.amazonaws.com" + with aws.S3_ENDPOINT.get_patcher(endpoint): + options = get_storage_options("s3://bucket/somewhere") + assert options == {"client_kwargs": {"endpoint_url": endpoint}} + + options = get_storage_options("/tmp/file") + assert options is None + + +cols = kwtypes(Name=str, Age=int) +subset_cols = kwtypes(Name=str) + + +@task +def t1( + df1: Annotated[pd.DataFrame, cols], df2: Annotated[pa.Table, cols] +) -> (Annotated[pd.DataFrame, subset_cols], Annotated[pa.Table, subset_cols]): + return df1, df2 + + +def test_structured_dataset_wf(): + pd_df = pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]}) + pa_df = pa.Table.from_pandas(pd_df) + + subset_pd_df = pd.DataFrame({"Name": ["Tom", "Joseph"]}) + subset_pa_df = pa.Table.from_pandas(subset_pd_df) + + df1, df2 = t1(df1=pd_df, df2=pa_df) + assert df1.equals(subset_pd_df) + assert df2.equals(subset_pa_df) diff --git a/plugins/flytekit-data-fsspec/tests/test_persist.py b/plugins/flytekit-data-fsspec/tests/test_persist.py index d2ac50a7f8..1114fbb91d 100644 --- a/plugins/flytekit-data-fsspec/tests/test_persist.py +++ b/plugins/flytekit-data-fsspec/tests/test_persist.py @@ -23,16 +23,24 @@ def test_s3_setup_args(): def test_get_protocol(): - assert FSSpecPersistence._get_protocol("s3://abc") == "s3" - assert FSSpecPersistence._get_protocol("/abc") == "file" - assert FSSpecPersistence._get_protocol("file://abc") == "file" - assert FSSpecPersistence._get_protocol("gs://abc") == "gs" - assert FSSpecPersistence._get_protocol("sftp://abc") == "sftp" - assert FSSpecPersistence._get_protocol("abfs://abc") == "abfs" + assert FSSpecPersistence.get_protocol("s3://abc") == "s3" + assert FSSpecPersistence.get_protocol("/abc") == "file" + assert FSSpecPersistence.get_protocol("file://abc") == "file" + assert FSSpecPersistence.get_protocol("gs://abc") == "gs" + assert FSSpecPersistence.get_protocol("sftp://abc") == "sftp" + assert FSSpecPersistence.get_protocol("abfs://abc") == "abfs" + + +def test_get_anonymous_filesystem(): + fs = FSSpecPersistence.get_anonymous_filesystem("/abc") + assert fs is None + fs = FSSpecPersistence.get_anonymous_filesystem("s3://abc") + assert fs is not None + assert fs.protocol == ["s3", "s3a"] def test_get_filesystem(): - fs = FSSpecPersistence._get_filesystem("/abc") + fs = FSSpecPersistence.get_filesystem("/abc") assert fs is not None assert isinstance(fs, LocalFileSystem) From cc3d739e86931e77618fd6154b88ec92c4a04f95 Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Tue, 22 Feb 2022 16:31:29 -0800 Subject: [PATCH 092/128] Delete unnecessary auth configuration (#858) Signed-off-by: Yee Hing Tong Signed-off-by: maximsmol --- flytekit/clients/raw.py | 199 +++++++++++++----- flytekit/clis/auth/auth.py | 16 +- flytekit/clis/auth/credentials.py | 56 +---- flytekit/clis/auth/discovery.py | 73 ------- flytekit/clis/sdk_in_container/basic_auth.py | 61 ------ flytekit/configuration/creds.py | 61 +----- flytekit/configuration/platform.py | 17 -- setup.py | 1 + tests/flytekit/unit/cli/auth/test_auth.py | 11 + .../unit/cli/auth/test_credentials.py | 39 ---- .../flytekit/unit/cli/auth/test_discovery.py | 66 ------ .../unit/cli/pyflyte/test_basic_auth.py | 32 --- tests/flytekit/unit/clients/test_raw.py | 145 ++++++++++--- 13 files changed, 301 insertions(+), 476 deletions(-) delete mode 100644 flytekit/clis/auth/discovery.py delete mode 100644 flytekit/clis/sdk_in_container/basic_auth.py delete mode 100644 tests/flytekit/unit/cli/auth/test_credentials.py delete mode 100644 tests/flytekit/unit/cli/auth/test_discovery.py delete mode 100644 tests/flytekit/unit/cli/pyflyte/test_basic_auth.py diff --git a/flytekit/clients/raw.py b/flytekit/clients/raw.py index ce58d8cc57..24d3e118b3 100644 --- a/flytekit/clients/raw.py +++ b/flytekit/clients/raw.py @@ -1,8 +1,15 @@ +from __future__ import annotations + +import base64 as _base64 +import logging as _logging import subprocess import time -from typing import List +from typing import Optional +import requests as _requests from flyteidl.service import admin_pb2_grpc as _admin_service +from flyteidl.service import auth_pb2 +from flyteidl.service import auth_pb2_grpc as auth_service from google.protobuf.json_format import MessageToJson as _MessageToJson from grpc import RpcError as _RpcError from grpc import StatusCode as _GrpcStatusCode @@ -11,49 +18,57 @@ from grpc import ssl_channel_credentials as _ssl_channel_credentials from flytekit.clis.auth import credentials as _credentials_access -from flytekit.clis.sdk_in_container import basic_auth as _basic_auth -from flytekit.configuration import creds as _creds_config -from flytekit.configuration.creds import _DEPRECATED_CLIENT_CREDENTIALS_SCOPE as _DEPRECATED_SCOPE +from flytekit.configuration import creds as creds_config +from flytekit.configuration.creds import CLIENT_CREDENTIALS_SECRET as _CREDENTIALS_SECRET from flytekit.configuration.creds import CLIENT_ID as _CLIENT_ID from flytekit.configuration.creds import COMMAND as _COMMAND -from flytekit.configuration.creds import DEPRECATED_OAUTH_SCOPES, SCOPES -from flytekit.configuration.platform import AUTH as _AUTH from flytekit.exceptions import user as _user_exceptions +from flytekit.exceptions.user import FlyteAuthenticationException from flytekit.loggers import cli_logger +_utf_8 = "utf-8" + -def _refresh_credentials_standard(flyte_client): +def _refresh_credentials_standard(flyte_client: RawSynchronousFlyteClient): """ This function is used when the configuration value for AUTH_MODE is set to 'standard'. This either fetches the existing access token or initiates the flow to request a valid access token and store it. :param flyte_client: RawSynchronousFlyteClient :return: """ - - client = _credentials_access.get_client(flyte_client.url) - if client.can_refresh_token: + authorization_header_key = flyte_client.public_client_config.authorization_metadata_key or None + if not flyte_client.oauth2_metadata or not flyte_client.public_client_config: + raise ValueError( + "Raw Flyte client attempting client credentials flow but no response from Admin detected. " + "Check your Admin server's .well-known endpoints to make sure they're working as expected." + ) + client = _credentials_access.get_client( + redirect_endpoint=flyte_client.public_client_config.redirect_uri, + client_id=flyte_client.public_client_config.client_id, + scopes=flyte_client.public_client_config.scopes, + auth_endpoint=flyte_client.oauth2_metadata.authorization_endpoint, + token_endpoint=flyte_client.oauth2_metadata.token_endpoint, + ) + if client.has_valid_credentials and not flyte_client.check_access_token(client.credentials.access_token): + # When Python starts up, if credentials have been stored in the keyring, then the AuthorizationClient + # will have read them into its _credentials field, but it won't be in the RawSynchronousFlyteClient's + # metadata field yet. Therefore, if there's a mismatch, copy it over. + flyte_client.set_access_token(client.credentials.access_token, authorization_header_key) + # However, after copying over credentials from the AuthorizationClient, we have to clear it to avoid the + # scenario where the stored credentials in the keyring are expired. If that's the case, then we only try + # them once (because client here is a singleton), and the next time, we'll do one of the two other conditions + # below. + client.clear() + return + elif client.can_refresh_token: client.refresh_access_token() + else: + client.start_authorization_flow() - flyte_client.set_access_token(client.credentials.access_token) - - -def _get_basic_flow_scopes() -> List[str]: - """ - Merge the scope value between the old scope config option and the new list option. - - :return: The scopes to use for basic auth flow. - """ - deprecated_single_scope = _DEPRECATED_SCOPE.get() - if deprecated_single_scope: - return [deprecated_single_scope] - scopes = DEPRECATED_OAUTH_SCOPES.get() or SCOPES.get() - if "openid" in scopes: - cli_logger.warning("Basic flow authentication should never use openid.") - - return scopes + flyte_client.set_access_token(client.credentials.access_token, authorization_header_key) -def _refresh_credentials_basic(flyte_client): +def _refresh_credentials_basic(flyte_client: RawSynchronousFlyteClient): """ This function is used by the _handle_rpc_error() decorator, depending on the AUTH_MODE config object. This handler is meant for SDK use-cases of auth (like pyflyte, or when users call SDK functions that require access to Admin, @@ -63,16 +78,24 @@ def _refresh_credentials_basic(flyte_client): :param flyte_client: RawSynchronousFlyteClient :return: """ - auth_endpoints = _credentials_access.get_authorization_endpoints(flyte_client.url) - token_endpoint = auth_endpoints.token_endpoint - client_secret = _basic_auth.get_secret() - cli_logger.debug( - "Basic authorization flow with client id {} scope {}".format(_CLIENT_ID.get(), _get_basic_flow_scopes()) - ) - authorization_header = _basic_auth.get_basic_authorization_header(_CLIENT_ID.get(), client_secret) - token, expires_in = _basic_auth.get_token(token_endpoint, authorization_header, _get_basic_flow_scopes()) + if not flyte_client.oauth2_metadata or not flyte_client.public_client_config: + raise ValueError( + "Raw Flyte client attempting client credentials flow but no response from Admin detected. " + "Check your Admin server's .well-known endpoints to make sure they're working as expected." + ) + + token_endpoint = flyte_client.oauth2_metadata.token_endpoint + scopes = creds_config.SCOPES.get() or flyte_client.public_client_config.scopes + scopes = ",".join(scopes) + + # Note that unlike the Pkce flow, the client ID does not come from Admin. + client_secret = get_secret() + cli_logger.debug("Basic authorization flow with client id {} scope {}".format(_CLIENT_ID.get(), scopes)) + authorization_header = get_basic_authorization_header(_CLIENT_ID.get(), client_secret) + token, expires_in = get_token(token_endpoint, authorization_header, scopes) cli_logger.info("Retrieved new token, expires in {}".format(expires_in)) - flyte_client.set_access_token(token) + authorization_header_key = flyte_client.public_client_config.authorization_metadata_key or None + flyte_client.set_access_token(token, authorization_header_key) def _refresh_credentials_from_command(flyte_client): @@ -101,7 +124,7 @@ def _refresh_credentials_noop(flyte_client): def _get_refresh_handler(auth_mode): if auth_mode == "standard": return _refresh_credentials_standard - elif auth_mode == "basic": + elif auth_mode == "basic" or auth_mode == "client_credentials": return _refresh_credentials_basic elif auth_mode == "external_process": return _refresh_credentials_from_command @@ -133,7 +156,7 @@ def handler(*args, **kwargs): # Exit the loop and wrap the authentication error. raise _user_exceptions.FlyteAuthenticationException(str(e)) cli_logger.error(f"Unauthenticated RPC error {e}, refreshing credentials and retrying\n") - refresh_handler_fn = _get_refresh_handler(_creds_config.AUTH_MODE.get()) + refresh_handler_fn = _get_refresh_handler(creds_config.AUTH_MODE.get()) refresh_handler_fn(args[0]) # There are two cases that we should throw error immediately # 1. Entity already exists when we register entity @@ -210,29 +233,57 @@ def __init__(self, url, insecure=False, credentials=None, options=None, root_cer options=list((options or {}).items()), ) self._stub = _admin_service.AdminServiceStub(self._channel) + self._auth_stub = auth_service.AuthMetadataServiceStub(self._channel) + try: + resp = self._auth_stub.GetPublicClientConfig(auth_pb2.PublicClientAuthConfigRequest()) + self._public_client_config = resp + except _RpcError: + cli_logger.debug("No public client auth config found, skipping.") + self._public_client_config = None + try: + resp = self._auth_stub.GetOAuth2Metadata(auth_pb2.OAuth2MetadataRequest()) + self._oauth2_metadata = resp + except _RpcError: + cli_logger.debug("No OAuth2 Metadata found, skipping.") + self._oauth2_metadata = None + + # metadata will hold the value of the token to send to the various endpoints. self._metadata = None - if _AUTH.get(): - self.force_auth_flow() + + @property + def public_client_config(self) -> Optional[auth_pb2.PublicClientAuthConfigResponse]: + return self._public_client_config + + @property + def oauth2_metadata(self) -> Optional[auth_pb2.OAuth2MetadataResponse]: + return self._oauth2_metadata @property def url(self) -> str: return self._url - def set_access_token(self, access_token): + def set_access_token(self, access_token: str, authorization_header_key: Optional[str] = "authorization"): # Always set the header to lower-case regardless of what the config is. The grpc libraries that Admin uses # to parse the metadata don't change the metadata, but they do automatically lower the key you're looking for. - authorization_metadata_key = _creds_config.AUTHORIZATION_METADATA_KEY.get().lower() - cli_logger.debug(f"Adding authorization header. Header name: {authorization_metadata_key}.") + cli_logger.debug(f"Adding authorization header. Header name: {authorization_header_key}.") self._metadata = [ ( - authorization_metadata_key, + authorization_header_key, f"Bearer {access_token}", ) ] - def force_auth_flow(self): - refresh_handler_fn = _get_refresh_handler(_creds_config.AUTH_MODE.get()) - refresh_handler_fn(self) + def check_access_token(self, access_token: str) -> bool: + """ + This checks to see if the given access token is the same as the one already stored in the client. The reason + this is useful is so that we can prevent unnecessary refreshing of tokens. + + :param access_token: The access token to check + :return: If no access token is stored, or if the stored token doesn't match, return False. + """ + if self._metadata is None: + return False + return access_token == self._metadata[0][1].replace("Bearer ", "") #################################################################################################################### # @@ -749,3 +800,55 @@ def list_matchable_attributes(self, matchable_attributes_list_request): # TODO: (P2) Implement the event endpoints in case there becomes a use-case for third-parties to submit events # through the client in Python. + + +def get_token(token_endpoint, authorization_header, scope): + """ + :param Text token_endpoint: + :param Text authorization_header: This is the value for the "Authorization" key. (eg 'Bearer abc123') + :param Text scope: + :rtype: (Text,Int) The first element is the access token retrieved from the IDP, the second is the expiration + in seconds + """ + headers = { + "Authorization": authorization_header, + "Cache-Control": "no-cache", + "Accept": "application/json", + "Content-Type": "application/x-www-form-urlencoded", + } + body = { + "grant_type": "client_credentials", + } + if scope is not None: + body["scope"] = scope + response = _requests.post(token_endpoint, data=body, headers=headers) + if response.status_code != 200: + _logging.error("Non-200 ({}) received from IDP: {}".format(response.status_code, response.text)) + raise FlyteAuthenticationException("Non-200 received from IDP") + + response = response.json() + return response["access_token"], response["expires_in"] + + +def get_secret(): + """ + This function will either read in the password from the file path given by the CLIENT_CREDENTIALS_SECRET_LOCATION + config object, or from the environment variable using the CLIENT_CREDENTIALS_SECRET config object. + :rtype: Text + """ + secret = _CREDENTIALS_SECRET.get() + if secret: + return secret + raise FlyteAuthenticationException("No secret could be found") + + +def get_basic_authorization_header(client_id, client_secret): + """ + This function transforms the client id and the client secret into a header that conforms with http basic auth. + It joins the id and the secret with a : then base64 encodes it, then adds the appropriate text. + :param Text client_id: + :param Text client_secret: + :rtype: Text + """ + concated = "{}:{}".format(client_id, client_secret) + return "Basic {}".format(_base64.b64encode(concated.encode(_utf_8)).decode(_utf_8)) diff --git a/flytekit/clis/auth/auth.py b/flytekit/clis/auth/auth.py index 79127c6748..2f27dfc72d 100644 --- a/flytekit/clis/auth/auth.py +++ b/flytekit/clis/auth/auth.py @@ -145,12 +145,11 @@ def __init__( scopes=None, client_id=None, redirect_uri=None, - client_secret=None, ): self._auth_endpoint = auth_endpoint self._token_endpoint = token_endpoint self._client_id = client_id - self._scopes = scopes + self._scopes = scopes or [] self._redirect_uri = redirect_uri self._code_verifier = _generate_code_verifier() code_challenge = _create_code_challenge(self._code_verifier) @@ -161,12 +160,11 @@ def __init__( self._refresh_token = None self._headers = {"content-type": "application/x-www-form-urlencoded"} self._expired = False - self._client_secret = client_secret self._params = { "client_id": client_id, # This must match the Client ID of the OAuth application. "response_type": "code", # Indicates the authorization code grant - "scope": " ".join(s.strip("' ") for s in scopes).strip( + "scope": " ".join(s.strip("' ") for s in self._scopes).strip( "[]'" ), # ensures that the /token endpoint returns an ID and refresh token # callback location where the user-agent will be directed to. @@ -239,12 +237,12 @@ def _initialize_credentials(self, auth_token_resp): raise ValueError('Expected "access_token" in response from oauth server') if "refresh_token" in response_body: self._refresh_token = response_body["refresh_token"] + _keyring.set_password( + _keyring_service_name, _keyring_refresh_token_storage_key, response_body["refresh_token"] + ) access_token = response_body["access_token"] - refresh_token = response_body["refresh_token"] - _keyring.set_password(_keyring_service_name, _keyring_access_token_storage_key, access_token) - _keyring.set_password(_keyring_service_name, _keyring_refresh_token_storage_key, refresh_token) self._credentials = Credentials(access_token=access_token) def request_access_token(self, auth_code): @@ -299,6 +297,10 @@ def credentials(self): """ return self._credentials + def clear(self): + self._credentials = None + self._refresh_token = None + @property def expired(self): """ diff --git a/flytekit/clis/auth/credentials.py b/flytekit/clis/auth/credentials.py index 45f51a482a..a8475c8dfc 100644 --- a/flytekit/clis/auth/credentials.py +++ b/flytekit/clis/auth/credentials.py @@ -1,56 +1,28 @@ -import urllib.parse as _urlparse +from typing import List -from flytekit.clis.auth.auth import AuthorizationClient as _AuthorizationClient -from flytekit.clis.auth.discovery import DiscoveryClient as _DiscoveryClient -from flytekit.configuration.creds import CLIENT_CREDENTIALS_SECRET as _CLIENT_SECRET -from flytekit.configuration.creds import CLIENT_ID as _CLIENT_ID -from flytekit.configuration.creds import DEPRECATED_OAUTH_SCOPES -from flytekit.configuration.creds import REDIRECT_URI as _REDIRECT_URI -from flytekit.configuration.creds import SCOPES -from flytekit.configuration.platform import HTTP_URL as _HTTP_URL -from flytekit.configuration.platform import INSECURE as _INSECURE -from flytekit.configuration.platform import URL as _URL +from flytekit.clis.auth.auth import AuthorizationClient from flytekit.loggers import auth_logger # Default, well known-URI string used for fetching JSON metadata. See https://tools.ietf.org/html/rfc8414#section-3. discovery_endpoint_path = "./.well-known/oauth-authorization-server" - -def _get_discovery_endpoint(http_config_val, platform_url_val, insecure_val): - if http_config_val: - scheme, netloc, path, _, _, _ = _urlparse.urlparse(http_config_val) - if not scheme: - scheme = "http" if insecure_val else "https" - else: # Use the main _URL config object effectively - scheme = "http" if insecure_val else "https" - netloc = platform_url_val - path = "" - - computed_endpoint = _urlparse.urlunparse((scheme, netloc, path, None, None, None)) - # The urljoin function needs a trailing slash in order to append things correctly. Also, having an extra slash - # at the end is okay, it just gets stripped out. - computed_endpoint = _urlparse.urljoin(computed_endpoint + "/", discovery_endpoint_path) - auth_logger.debug(f"Using {computed_endpoint} as discovery endpoint") - return computed_endpoint - - # Lazy initialized authorization client singleton _authorization_client = None -def get_client(flyte_client_url): +def get_client( + redirect_endpoint: str, client_id: str, scopes: List[str], auth_endpoint: str, token_endpoint: str +) -> AuthorizationClient: global _authorization_client if _authorization_client is not None and not _authorization_client.expired: return _authorization_client - authorization_endpoints = get_authorization_endpoints(flyte_client_url) - _authorization_client = _AuthorizationClient( - redirect_uri=_REDIRECT_URI.get(), - client_id=_CLIENT_ID.get(), - scopes=DEPRECATED_OAUTH_SCOPES.get() or SCOPES.get(), - auth_endpoint=authorization_endpoints.auth_endpoint, - token_endpoint=authorization_endpoints.token_endpoint, - client_secret=_CLIENT_SECRET.get(), + _authorization_client = AuthorizationClient( + redirect_uri=redirect_endpoint, + client_id=client_id, + scopes=scopes, + auth_endpoint=auth_endpoint, + token_endpoint=token_endpoint, ) auth_logger.debug(f"Created oauth client with redirect {_authorization_client}") @@ -59,9 +31,3 @@ def get_client(flyte_client_url): _authorization_client.start_authorization_flow() return _authorization_client - - -def get_authorization_endpoints(flyte_client_url): - discovery_endpoint = _get_discovery_endpoint(_HTTP_URL.get(), flyte_client_url or _URL.get(), _INSECURE.get()) - discovery_client = _DiscoveryClient(discovery_url=discovery_endpoint) - return discovery_client.get_authorization_endpoints() diff --git a/flytekit/clis/auth/discovery.py b/flytekit/clis/auth/discovery.py deleted file mode 100644 index 5134eab974..0000000000 --- a/flytekit/clis/auth/discovery.py +++ /dev/null @@ -1,73 +0,0 @@ -import logging - -import requests as _requests - -# These response keys are defined in https://tools.ietf.org/id/draft-ietf-oauth-discovery-08.html. -_authorization_endpoint_key = "authorization_endpoint" -_token_endpoint_key = "token_endpoint" - - -class AuthorizationEndpoints(object): - """ - A simple wrapper around commonly discovered endpoints used for the PKCE auth flow. - """ - - def __init__(self, auth_endpoint=None, token_endpoint=None): - self._auth_endpoint = auth_endpoint - self._token_endpoint = token_endpoint - - @property - def auth_endpoint(self): - return self._auth_endpoint - - @property - def token_endpoint(self): - return self._token_endpoint - - -class DiscoveryClient(object): - """ - Discovers well known OpenID configuration and parses out authorization endpoints required for initiating the PKCE - auth flow. - """ - - def __init__(self, discovery_url=None): - logging.debug("Initializing discovery client with {}".format(discovery_url)) - self._discovery_url = discovery_url - self._authorization_endpoints = None - - @property - def authorization_endpoints(self): - """ - :rtype: flytekit.clis.auth.discovery.AuthorizationEndpoints: - """ - return self._authorization_endpoints - - def get_authorization_endpoints(self): - if self.authorization_endpoints is not None: - return self.authorization_endpoints - resp = _requests.get( - url=self._discovery_url, - ) - - response_body = resp.json() - - authorization_endpoint = response_body[_authorization_endpoint_key] - token_endpoint = response_body[_token_endpoint_key] - - if authorization_endpoint is None: - raise ValueError("Unable to discover authorization endpoint") - - if token_endpoint is None: - raise ValueError("Unable to discover token endpoint") - - if authorization_endpoint.startswith("/"): - authorization_endpoint = _requests.compat.urljoin(self._discovery_url, authorization_endpoint) - - if token_endpoint.startswith("/"): - token_endpoint = _requests.compat.urljoin(self._discovery_url, token_endpoint) - - self._authorization_endpoints = AuthorizationEndpoints( - auth_endpoint=authorization_endpoint, token_endpoint=token_endpoint - ) - return self.authorization_endpoints diff --git a/flytekit/clis/sdk_in_container/basic_auth.py b/flytekit/clis/sdk_in_container/basic_auth.py deleted file mode 100644 index 92612595ad..0000000000 --- a/flytekit/clis/sdk_in_container/basic_auth.py +++ /dev/null @@ -1,61 +0,0 @@ -import base64 as _base64 -import logging as _logging - -import requests as _requests - -from flytekit.configuration.creds import CLIENT_CREDENTIALS_SECRET as _CREDENTIALS_SECRET -from flytekit.exceptions.user import FlyteAuthenticationException - -_utf_8 = "utf-8" - - -def get_secret(): - """ - This function will either read in the password from the file path given by the CLIENT_CREDENTIALS_SECRET_LOCATION - config object, or from the environment variable using the CLIENT_CREDENTIALS_SECRET config object. - :rtype: Text - """ - secret = _CREDENTIALS_SECRET.get() - if secret: - return secret - raise FlyteAuthenticationException("No secret could be found") - - -def get_basic_authorization_header(client_id, client_secret): - """ - This function transforms the client id and the client secret into a header that conforms with http basic auth. - It joins the id and the secret with a : then base64 encodes it, then adds the appropriate text. - :param Text client_id: - :param Text client_secret: - :rtype: Text - """ - concated = "{}:{}".format(client_id, client_secret) - return "Basic {}".format(_base64.b64encode(concated.encode(_utf_8)).decode(_utf_8)) - - -def get_token(token_endpoint, authorization_header, scope): - """ - :param Text token_endpoint: - :param Text authorization_header: This is the value for the "Authorization" key. (eg 'Bearer abc123') - :param Text scope: - :rtype: (Text,Int) The first element is the access token retrieved from the IDP, the second is the expiration - in seconds - """ - headers = { - "Authorization": authorization_header, - "Cache-Control": "no-cache", - "Accept": "application/json", - "Content-Type": "application/x-www-form-urlencoded", - } - body = { - "grant_type": "client_credentials", - } - if scope is not None: - body["scope"] = scope - response = _requests.post(token_endpoint, data=body, headers=headers) - if response.status_code != 200: - _logging.error("Non-200 ({}) received from IDP: {}".format(response.status_code, response.text)) - raise FlyteAuthenticationException("Non-200 received from IDP") - - response = response.json() - return response["access_token"], response["expires_in"] diff --git a/flytekit/configuration/creds.py b/flytekit/configuration/creds.py index 6acfcd9129..9f11ac2d2e 100644 --- a/flytekit/configuration/creds.py +++ b/flytekit/configuration/creds.py @@ -1,10 +1,5 @@ -from warnings import warn - from flytekit.configuration import common as _config_common -deprecated_names = ["CLIENT_CREDENTIALS_SCOPE"] - - COMMAND = _config_common.FlyteStringListConfigurationEntry("credentials", "command", default=None) """ This command is executed to return a token using an external process. @@ -16,37 +11,6 @@ More details here: https://www.oauth.com/oauth2-servers/client-registration/client-id-secret/. """ -REDIRECT_URI = _config_common.FlyteStringConfigurationEntry( - "credentials", "redirect_uri", default="http://localhost:12345/callback" -) -""" -This is the callback uri registered with the app which handles authorization for a Flyte deployment. -Please note the hardcoded port number. Ideally we would not do this, but some IDPs do not allow wildcards for -the URL, which means we have to use the same port every time. This is the only reason this is a configuration option, -otherwise, we'd just hardcode the callback path as a constant. -FYI, to see if a given port is already in use, run `sudo lsof -i :` if on a Linux system. -More details here: https://www.oauth.com/oauth2-servers/redirect-uris/. -""" - -SCOPES = _config_common.FlyteStringListConfigurationEntry("credentials", "scopes", default=["openid"]) -""" -This controls the list of scopes to request from the authorization server. -""" - -DEPRECATED_OAUTH_SCOPES = _config_common.FlyteStringListConfigurationEntry("credentials", "oauth_scopes", default=None) -""" -This controls the list of scopes to request from the authorization server. -Deprecated - please use the SCOPES variable. -""" - -AUTHORIZATION_METADATA_KEY = _config_common.FlyteStringConfigurationEntry( - "credentials", "authorization_metadata_key", default="authorization" -) -""" -The authorization metadata key used for passing access tokens in gRPC requests. -Traditionally this value is 'authorization' however it is made configurable. -""" - CLIENT_CREDENTIALS_SECRET = _config_common.FlyteStringConfigurationEntry("credentials", "client_secret", default=None) """ Used for basic auth, which is automatically called during pyflyte. This will allow the Flyte engine to read the @@ -54,33 +18,14 @@ secret as a file is impossible. """ -_DEPRECATED_CLIENT_CREDENTIALS_SCOPE = _config_common.FlyteStringConfigurationEntry( - "credentials", "scope", default=None -) -""" -Used for basic auth, which is automatically called during pyflyte. This is the scope that will be requested. Because -there is no user explicitly in this auth flow, certain IDPs require a custom scope for basic auth in the configuration -of the authorization server. - -Deprecated - please use the OAUTH_SCOPES list variable instead. In the basic flow scenario, flytekit will expect a list -with at least one element. The first element will be used. If list has more than one element a warning will be logged. -Config files with both this option, and the OAUTH_SCOPES, will use this one. -""" +SCOPES = _config_common.FlyteStringListConfigurationEntry("credentials", "scopes", default=[]) AUTH_MODE = _config_common.FlyteStringConfigurationEntry("credentials", "auth_mode", default="standard") """ The auth mode defines the behavior used to request and refresh credentials. The currently supported modes include: - 'standard' This uses the pkce-enhanced authorization code flow by opening a browser window to initiate credentials access. -- 'basic' This uses cert-based auth in which the end user enters his/her username and password and public key encryption - is used to facilitate authentication. +- 'basic' or 'client_credentials' This uses cert-based auth in which the end user enters a client id and a client + secret and public key encryption is used to facilitate authentication. - None: No auth will be attempted. """ - - -# https://www.python.org/dev/peps/pep-0562/ -def __getattr__(name): - if name in deprecated_names: - warn(f"{name} is deprecated", DeprecationWarning) - return globals()[f"_DEPRECATED_{name}"] - raise AttributeError(f"module {__name__} has no attribute {name}") diff --git a/flytekit/configuration/platform.py b/flytekit/configuration/platform.py index eecbeda162..ee8bfeb895 100644 --- a/flytekit/configuration/platform.py +++ b/flytekit/configuration/platform.py @@ -1,21 +1,4 @@ from flytekit.configuration import common as _config_common URL = _config_common.FlyteStringConfigurationEntry("platform", "url") - -HTTP_URL = _config_common.FlyteStringConfigurationEntry("platform", "http_url", default=None) -""" -If not starting with either http or https, this setting should begin with // as per the urlparse library and -https://tools.ietf.org/html/rfc1808.html, otherwise the netloc will not be properly parsed. - -Currently the only use-case for this configuration setting is for Auth discovery. This setting supports the case where -Flyte Admin's gRPC and HTTP points are deployed on different ports. -""" - INSECURE = _config_common.FlyteBoolConfigurationEntry("platform", "insecure", default=False) - -AUTH = _config_common.FlyteBoolConfigurationEntry("platform", "auth", default=False) -""" -This config setting should not normally be filled in. Whether or not an admin server requires authentication should be -something published by the admin server itself (typically by returning a 401). However, to help with migration, this -config object is here to force the SDK to attempt the auth flow even without prompting by Admin. -""" diff --git a/setup.py b/setup.py index 438fd9a1a6..7b08234cfd 100644 --- a/setup.py +++ b/setup.py @@ -46,6 +46,7 @@ "python-dateutil>=2.1", "grpcio>=1.3.0,<2.0", "protobuf>=3.6.1,<4", + "protoc_gen_swagger", "python-json-logger>=2.0.0", "pytimeparse>=1.1.8,<2.0.0", "pytz", diff --git a/tests/flytekit/unit/cli/auth/test_auth.py b/tests/flytekit/unit/cli/auth/test_auth.py index 2deecaefaa..1bd38d3b39 100644 --- a/tests/flytekit/unit/cli/auth/test_auth.py +++ b/tests/flytekit/unit/cli/auth/test_auth.py @@ -1,6 +1,8 @@ import re from multiprocessing import Queue as _Queue +from mock import patch + from flytekit.clis.auth import auth as _auth try: # Python 3 @@ -33,3 +35,12 @@ def test_oauth_http_server(): server.handle_authorization_code(test_auth_code) auth_code = queue.get() assert test_auth_code == auth_code + + +@patch("flytekit.clis.auth.auth._keyring.get_password") +def test_clear(mock_get_password): + mock_get_password.return_value = "token" + ac = _auth.AuthorizationClient() + ac.clear() + assert ac.credentials is None + assert not ac.can_refresh_token diff --git a/tests/flytekit/unit/cli/auth/test_credentials.py b/tests/flytekit/unit/cli/auth/test_credentials.py deleted file mode 100644 index f1ae57a016..0000000000 --- a/tests/flytekit/unit/cli/auth/test_credentials.py +++ /dev/null @@ -1,39 +0,0 @@ -from flytekit.clis.auth import credentials as _credentials - - -def test_get_discovery_endpoint(): - endpoint = _credentials._get_discovery_endpoint("//localhost:8088", "localhost:8089", True) - assert endpoint == "http://localhost:8088/.well-known/oauth-authorization-server" - - endpoint = _credentials._get_discovery_endpoint("//localhost:8088", "localhost:8089", False) - assert endpoint == "https://localhost:8088/.well-known/oauth-authorization-server" - - endpoint = _credentials._get_discovery_endpoint("//localhost:8088/path", "localhost:8089", True) - assert endpoint == "http://localhost:8088/path/.well-known/oauth-authorization-server" - - endpoint = _credentials._get_discovery_endpoint("//localhost:8088/path", "localhost:8089", False) - assert endpoint == "https://localhost:8088/path/.well-known/oauth-authorization-server" - - endpoint = _credentials._get_discovery_endpoint("//flyte.corp.com", "localhost:8089", False) - assert endpoint == "https://flyte.corp.com/.well-known/oauth-authorization-server" - - endpoint = _credentials._get_discovery_endpoint("//flyte.corp.com/path", "localhost:8089", False) - assert endpoint == "https://flyte.corp.com/path/.well-known/oauth-authorization-server" - - endpoint = _credentials._get_discovery_endpoint(None, "localhost:8089", True) - assert endpoint == "http://localhost:8089/.well-known/oauth-authorization-server" - - endpoint = _credentials._get_discovery_endpoint(None, "localhost:8089", False) - assert endpoint == "https://localhost:8089/.well-known/oauth-authorization-server" - - endpoint = _credentials._get_discovery_endpoint(None, "flyte.corp.com", True) - assert endpoint == "http://flyte.corp.com/.well-known/oauth-authorization-server" - - endpoint = _credentials._get_discovery_endpoint(None, "flyte.corp.com", False) - assert endpoint == "https://flyte.corp.com/.well-known/oauth-authorization-server" - - endpoint = _credentials._get_discovery_endpoint(None, "localhost:8089", True) - assert endpoint == "http://localhost:8089/.well-known/oauth-authorization-server" - - endpoint = _credentials._get_discovery_endpoint(None, "localhost:8089", False) - assert endpoint == "https://localhost:8089/.well-known/oauth-authorization-server" diff --git a/tests/flytekit/unit/cli/auth/test_discovery.py b/tests/flytekit/unit/cli/auth/test_discovery.py deleted file mode 100644 index c75427f35d..0000000000 --- a/tests/flytekit/unit/cli/auth/test_discovery.py +++ /dev/null @@ -1,66 +0,0 @@ -import pytest -import responses - -from flytekit.clis.auth import discovery as _discovery - - -@responses.activate -def test_get_authorization_endpoints(): - discovery_url = "http://flyte-admin.com/discovery" - - auth_endpoint = "http://flyte-admin.com/authorization" - token_endpoint = "http://flyte-admin.com/token" - responses.add( - responses.GET, - discovery_url, - json={"authorization_endpoint": auth_endpoint, "token_endpoint": token_endpoint}, - ) - - discovery_client = _discovery.DiscoveryClient(discovery_url=discovery_url) - assert discovery_client.get_authorization_endpoints().auth_endpoint == auth_endpoint - assert discovery_client.get_authorization_endpoints().token_endpoint == token_endpoint - - -@responses.activate -def test_get_authorization_endpoints_relative(): - discovery_url = "http://flyte-admin.com/discovery" - - auth_endpoint = "/authorization" - token_endpoint = "/token" - responses.add( - responses.GET, - discovery_url, - json={"authorization_endpoint": auth_endpoint, "token_endpoint": token_endpoint}, - ) - - discovery_client = _discovery.DiscoveryClient(discovery_url=discovery_url) - assert discovery_client.get_authorization_endpoints().auth_endpoint == "http://flyte-admin.com/authorization" - assert discovery_client.get_authorization_endpoints().token_endpoint == "http://flyte-admin.com/token" - - -@responses.activate -def test_get_authorization_endpoints_missing_authorization_endpoint(): - discovery_url = "http://flyte-admin.com/discovery" - responses.add( - responses.GET, - discovery_url, - json={"token_endpoint": "http://flyte-admin.com/token"}, - ) - - discovery_client = _discovery.DiscoveryClient(discovery_url=discovery_url) - with pytest.raises(Exception): - discovery_client.get_authorization_endpoints() - - -@responses.activate -def test_get_authorization_endpoints_missing_token_endpoint(): - discovery_url = "http://flyte-admin.com/discovery" - responses.add( - responses.GET, - discovery_url, - json={"authorization_endpoint": "http://flyte-admin.com/authorization"}, - ) - - discovery_client = _discovery.DiscoveryClient(discovery_url=discovery_url) - with pytest.raises(Exception): - discovery_client.get_authorization_endpoints() diff --git a/tests/flytekit/unit/cli/pyflyte/test_basic_auth.py b/tests/flytekit/unit/cli/pyflyte/test_basic_auth.py deleted file mode 100644 index d18f21dfa5..0000000000 --- a/tests/flytekit/unit/cli/pyflyte/test_basic_auth.py +++ /dev/null @@ -1,32 +0,0 @@ -import json - -from mock import MagicMock, patch - -from flytekit.clis.flyte_cli.main import _welcome_message -from flytekit.clis.sdk_in_container import basic_auth -from flytekit.configuration.creds import CLIENT_CREDENTIALS_SECRET as _CREDENTIALS_SECRET - -_welcome_message() - - -def test_get_secret(): - import os - - os.environ[_CREDENTIALS_SECRET.env_var] = "abc" - assert basic_auth.get_secret() == "abc" - - -def test_get_basic_authorization_header(): - header = basic_auth.get_basic_authorization_header("client_id", "abc") - assert header == "Basic Y2xpZW50X2lkOmFiYw==" - - -@patch("flytekit.clis.sdk_in_container.basic_auth._requests") -def test_get_token(mock_requests): - response = MagicMock() - response.status_code = 200 - response.json.return_value = json.loads("""{"access_token": "abc", "expires_in": 60}""") - mock_requests.post.return_value = response - access, expiration = basic_auth.get_token("https://corp.idp.net", "abc123", "my_scope") - assert access == "abc" - assert expiration == 60 diff --git a/tests/flytekit/unit/clients/test_raw.py b/tests/flytekit/unit/clients/test_raw.py index c8e56811d4..86f5a04a6b 100644 --- a/tests/flytekit/unit/clients/test_raw.py +++ b/tests/flytekit/unit/clients/test_raw.py @@ -3,44 +3,59 @@ from subprocess import CompletedProcess import mock +import pytest from flyteidl.admin import project_pb2 as _project_pb2 +from flyteidl.service import auth_pb2 +from mock import MagicMock, patch from flytekit.clients.raw import RawSynchronousFlyteClient as _RawSynchronousFlyteClient -from flytekit.clients.raw import _get_basic_flow_scopes, _refresh_credentials_basic, _refresh_credentials_from_command -from flytekit.clis.auth.discovery import AuthorizationEndpoints as _AuthorizationEndpoints -from flytekit.configuration import TemporaryConfiguration +from flytekit.clients.raw import ( + _get_refresh_handler, + _refresh_credentials_basic, + _refresh_credentials_from_command, + _refresh_credentials_standard, + get_basic_authorization_header, + get_secret, + get_token, +) from flytekit.configuration.creds import CLIENT_CREDENTIALS_SECRET as _CREDENTIALS_SECRET -@mock.patch("flytekit.clients.raw.RawSynchronousFlyteClient.force_auth_flow") +def get_admin_stub_mock() -> mock.MagicMock: + auth_stub_mock = mock.MagicMock() + auth_stub_mock.GetPublicClientConfig.return_value = auth_pb2.PublicClientAuthConfigResponse( + client_id="flytectl", + redirect_uri="http://localhost:53593/callback", + scopes=["offline", "all"], + authorization_metadata_key="flyte-authorization", + ) + auth_stub_mock.GetOAuth2Metadata.return_value = auth_pb2.OAuth2MetadataResponse( + issuer="https://your.domain.io", + authorization_endpoint="https://your.domain.io/oauth2/authorize", + token_endpoint="https://your.domain.io/oauth2/token", + response_types_supported=["code", "token", "code token"], + scopes_supported=["all"], + token_endpoint_auth_methods_supported=["client_secret_basic"], + jwks_uri="https://your.domain.io/oauth2/jwks", + code_challenge_methods_supported=["S256"], + grant_types_supported=["client_credentials", "refresh_token", "authorization_code"], + ) + return auth_stub_mock + + +@mock.patch("flytekit.clients.raw.auth_service") @mock.patch("flytekit.clients.raw._admin_service") @mock.patch("flytekit.clients.raw._insecure_channel") @mock.patch("flytekit.clients.raw._secure_channel") -def test_client_set_token(mock_secure_channel, mock_channel, mock_admin, mock_force): - mock_force.return_value = True +def test_client_set_token(mock_secure_channel, mock_channel, mock_admin, mock_admin_auth): mock_secure_channel.return_value = True mock_channel.return_value = True mock_admin.AdminServiceStub.return_value = True + mock_admin_auth.AuthMetadataServiceStub.return_value = get_admin_stub_mock() client = _RawSynchronousFlyteClient(url="a.b.com", insecure=True) client.set_access_token("abc") assert client._metadata[0][1] == "Bearer abc" - - -@mock.patch("flytekit.clis.sdk_in_container.basic_auth._requests") -@mock.patch("flytekit.clients.raw._credentials_access") -def test_refresh_credentials_basic(mock_credentials_access, mock_requests): - mock_credentials_access.get_authorization_endpoints.return_value = _AuthorizationEndpoints("auth", "token") - response = mock.MagicMock() - response.status_code = 200 - response.json.return_value = json.loads("""{"access_token": "abc", "expires_in": 60}""") - mock_requests.post.return_value = response - os.environ[_CREDENTIALS_SECRET.env_var] = "asdf12345" - - mock_client = mock.MagicMock() - mock_client.url.return_value = "flyte.localhost" - _refresh_credentials_basic(mock_client) - mock_client.set_access_token.assert_called_with("abc") - mock_credentials_access.get_authorization_endpoints.assert_called_with(mock_client.url) + assert client.check_access_token("abc") @mock.patch("flytekit.configuration.creds.COMMAND.get") @@ -59,6 +74,57 @@ def test_refresh_credentials_from_command(mock_call_to_external_process, mock_co mock_client.set_access_token.assert_called_with(token) +@mock.patch("flytekit.configuration.creds.SCOPES.get") +@mock.patch("flytekit.clients.raw.get_secret") +@mock.patch("flytekit.clients.raw.get_basic_authorization_header") +@mock.patch("flytekit.clients.raw.get_token") +@mock.patch("flytekit.clients.raw.auth_service") +@mock.patch("flytekit.clients.raw._admin_service") +@mock.patch("flytekit.clients.raw._insecure_channel") +@mock.patch("flytekit.clients.raw._secure_channel") +def test_refresh_client_credentials_aka_basic( + mock_secure_channel, + mock_channel, + mock_admin, + mock_admin_auth, + mock_get_token, + mock_get_basic_header, + mock_secret, + mock_scopes, +): + mock_secret.return_value = "sosecret" + mock_scopes.return_value = ["a", "b", "c", "d"] + mock_secure_channel.return_value = True + mock_channel.return_value = True + mock_admin.AdminServiceStub.return_value = True + mock_get_basic_header.return_value = "Basic 123" + mock_get_token.return_value = ("token1", 1234567) + + mock_admin_auth.AuthMetadataServiceStub.return_value = get_admin_stub_mock() + client = _RawSynchronousFlyteClient(url="a.b.com", insecure=True) + client._metadata = None + assert not client.check_access_token("fdsa") + _refresh_credentials_basic(client) + + # Scopes from configuration take precendence. + mock_get_token.assert_called_once_with("https://your.domain.io/oauth2/token", "Basic 123", "a,b,c,d") + + client.set_access_token("token") + assert client._metadata[0][0] == "authorization" + + +def test_raises(): + mm = MagicMock() + mm.public_client_config = None + with pytest.raises(ValueError): + _refresh_credentials_basic(mm) + + mm = MagicMock() + mm.oauth2_metadata = None + with pytest.raises(ValueError): + _refresh_credentials_basic(mm) + + @mock.patch("flytekit.clients.raw._admin_service") @mock.patch("flytekit.clients.raw._insecure_channel") def test_update_project(mock_channel, mock_admin): @@ -77,12 +143,31 @@ def test_list_projects_paginated(mock_channel, mock_admin): mock_admin.AdminServiceStub().ListProjects.assert_called_with(project_list_request, metadata=None) -def test_scope_deprecation(): - with TemporaryConfiguration(os.path.join(os.path.dirname(__file__), "auth_deprecation.config")): - assert _get_basic_flow_scopes() == ["custom_basic"] +def test_get_secret(): + os.environ[_CREDENTIALS_SECRET.env_var] = "abc" + assert get_secret() == "abc" + + +def test_get_basic_authorization_header(): + header = get_basic_authorization_header("client_id", "abc") + assert header == "Basic Y2xpZW50X2lkOmFiYw==" - with TemporaryConfiguration(os.path.join(os.path.dirname(__file__), "auth_deprecation2.config")): - assert _get_basic_flow_scopes() == ["custom_basic", "other_scope", "profile"] - with TemporaryConfiguration(os.path.join(os.path.dirname(__file__), "auth_deprecation3.config")): - assert _get_basic_flow_scopes() == ["custom_basic"] +@patch("flytekit.clients.raw._requests") +def test_get_token(mock_requests): + response = MagicMock() + response.status_code = 200 + response.json.return_value = json.loads("""{"access_token": "abc", "expires_in": 60}""") + mock_requests.post.return_value = response + access, expiration = get_token("https://corp.idp.net", "abc123", "my_scope") + assert access == "abc" + assert expiration == 60 + + +def test_get_refresh_handler(): + cc = _get_refresh_handler("client_credentials") + basic = _get_refresh_handler("basic") + assert basic is cc + assert basic is _refresh_credentials_basic + standard = _get_refresh_handler("standard") + assert standard is _refresh_credentials_standard From 96ad23ffafbdfecf714c83471ab24afa2229e0a1 Mon Sep 17 00:00:00 2001 From: SmritiSatyanV <94349093+SmritiSatyanV@users.noreply.github.com> Date: Wed, 23 Feb 2022 10:25:59 +0530 Subject: [PATCH 093/128] Fixed format alias in Flytekit docs (#844) * Fixed format alias Fixed docs for file format alias that weren't rendered properly. A warning popped up stating 'ignore' is deprecated, and to use 'ignore_paths' instead. Signed-off-by: SmritiSatyanV --- .github/workflows/pythonbuild.yml | 5 +++ codecov.yml | 2 +- docs/source/_templates/file_types.rst | 39 ++++++++++++++++++ docs/source/contributing.rst | 16 ++++---- docs/source/design/authoring.rst | 34 ++++++++-------- docs/source/design/clis.rst | 2 +- docs/source/design/control_plane.rst | 14 +++---- docs/source/design/execution.rst | 10 ++--- docs/source/design/index.rst | 4 +- docs/source/design/models.rst | 4 +- docs/source/index.rst | 8 ++-- flytekit/types/directory/__init__.py | 7 ++-- flytekit/types/file/__init__.py | 58 +++++++++------------------ flytekit/types/file/file.py | 16 ++++---- flytekit/types/schema/types.py | 4 +- 15 files changed, 124 insertions(+), 99 deletions(-) create mode 100644 docs/source/_templates/file_types.rst diff --git a/.github/workflows/pythonbuild.yml b/.github/workflows/pythonbuild.yml index f1cfa4977f..970673e090 100644 --- a/.github/workflows/pythonbuild.yml +++ b/.github/workflows/pythonbuild.yml @@ -138,7 +138,12 @@ jobs: - name: ShellCheck uses: ludeeus/action-shellcheck@master with: +<<<<<<< HEAD ignore: boilerplate +======= + ignore_paths: + boilerplate +>>>>>>> 0d034c15 (Fixed format alias in Flytekit docs (#844)) docs: runs-on: ubuntu-latest diff --git a/codecov.yml b/codecov.yml index 45721a38f7..89ec57f646 100644 --- a/codecov.yml +++ b/codecov.yml @@ -1,4 +1,4 @@ -ignore: +ignore_paths: - "flytekit/bin" - "test_*.py" - "flytekit/__init__.py" diff --git a/docs/source/_templates/file_types.rst b/docs/source/_templates/file_types.rst new file mode 100644 index 0000000000..e7629ea363 --- /dev/null +++ b/docs/source/_templates/file_types.rst @@ -0,0 +1,39 @@ +{{ fullname | escape | underline}} + +.. currentmodule:: {{ module }} + +{% if objname == 'FlyteFile' %} + +.. autoclass:: {{ objname }} + + {% block methods %} + {% if methods %} + + .. rubric:: {{ _('Methods') }} + {% for item in methods %} + + {% if item != '__init__' %} + .. automethod:: {{ item }} + {% endif %} + + {%- endfor %} + {% endif %} + {% endblock %} + + {% block attributes %} + {% if attributes %} + + .. rubric:: {{ _('Attributes') }} + {% for item in attributes %} + .. autoattribute:: {{ item }} + {%- endfor %} + + {% endif %} + {% endblock %} + + +{% else %} + +.. autodata:: {{ objname }} + +{% endif %} diff --git a/docs/source/contributing.rst b/docs/source/contributing.rst index 2cf40e7a94..1a69f67c8a 100644 --- a/docs/source/contributing.rst +++ b/docs/source/contributing.rst @@ -12,9 +12,9 @@ Please also take some time to read through the :std:ref:`design guides ` 📜 Quick Background ******************* -The first version of the flytekit library was written circa 2017, before mypy typing was mainstream, and -targeted Python 2. That legacy code will be fully deprecated and removed in 2022 but because there are still -users of flytekit that rely on that legacy api, you'll see 2 separate and distinct code paths within this repo. +The first version of the Flytekit library was written circa 2017, before mypy typing was mainstream, and +targeted Python 2. That legacy code will be fully deprecated and removed in 2022. Since there are still +users of Flytekit that rely on that legacy API, you'll see 2 separate and distinct code paths within this repo. Users and contributors should ignore the legacy sections. Below is a listing of the most important packages that comprise the new API: @@ -25,7 +25,7 @@ comprise the new API: - ``flytekit/extend`` This is the future home of extension points, and currently serves as the raw documentation for extensions. - ``flytekit/extras`` - This contains code that we want bundled with flytekit but not everyone may find useful (for example AWS and GCP + This contains code that we want bundled with Flytekit but not everyone may find useful (for example AWS and GCP specific logic). - ``flytekit/remote`` This implements the interface to interact with the Flyte service. Think of the code here as the Python-object version of Console. @@ -39,7 +39,7 @@ comprise the new API: - ``flytekit/bin/entrypoint.py`` The run time entrypoint for flytekit. When a task kicks off, this is where the click command goes. - ``flytekit/clis`` - This is the home for the clis. + This is the home for the CLIs. - ``flytekit/configuration`` This holds all the configuration objects, but dependency on configuration should be carefully considered as it makes compiled Flyte tasks and workflows less portable (i.e. if you run ``pyflyte package`` can someone else use @@ -77,16 +77,16 @@ We recommend using a virtual environment to develop Flytekit. Inside the top lev Install `shellcheck `__ for linting shell scripts. .. note:: - It's important to maintain separate virtualenvs for flytekit *development* and flytekit *use*. The reason is that installing a Python + It's important to maintain separate virtualenvs for flytekit *development* and Flytekit *use*. The reason is that installing a Python library in editable mode will link it to your source code. That is, the behavior will change as you work on the code, check out different branches, etc. -This will install flytekit dependencies and also install flytekit itself in editable mode. This basically links your virtual Python's ``site-packages`` with your local repo folder, allowing your local changes to take effect when the same Python interpreter runs ``import flytekit``. +This will install Flytekit dependencies and also install Flytekit itself in editable mode. This basically links your virtual Python's ``site-packages`` with your local repo folder, allowing your local changes to take effect when the same Python interpreter runs ``import flytekit``. Plugin Development ================== -As discussed in the design component, Flytekit plugins currently live in this flytekit repo, but under a different top level folder ``plugins``. +As discussed in the design component, Flytekit plugins currently live in this Flytekit repo, but under a different top level folder ``plugins``. In the future, this will be separated out into a different repo. These plugins follow a `microlib `__ structure, which will persist even if we move repos. :: source ~/.virtualenvs/flytekit/bin/activate diff --git a/docs/source/design/authoring.rst b/docs/source/design/authoring.rst index 2b1d0ea705..aaa8f461cf 100644 --- a/docs/source/design/authoring.rst +++ b/docs/source/design/authoring.rst @@ -4,26 +4,26 @@ Authoring Structure ############################ -Enabling users to write tasks and workflows is the core feature of flytekit, it is why it exists. This document goes over how some of the internals work. +Enabling users to write tasks and workflows is the core feature of Flytekit, it is why it exists. This document goes over how some of the internals work. ************* Background ************* -Please see the `design doc `__. +Please refer `design doc `__. ********************* Types and Type Engine ********************* -Flyte has its own type system, which is codified `in the IDL `__. Python of course has its own typing system, even though it's a dynamic language, and is mostly explained in `PEP 484 `_. In order to work properly, flytekit needs to be able to convert between the two. +Flyte has its own type system, which is codified `in the IDL `__. Python, of course, has its own typing system, even though it's a dynamic language, and is mostly explained in `PEP 484 `_. In order to work properly, Flytekit needs to be able to convert between the two. Type Engine ============= -The primary way this happens is through the :py:class:`flytekit.extend.TypeEngine`. This engine works by invoking a series of :py:class:`TypeTransformers `. Each transformer is responsible for providing the functionality that the engine needs for a given native Python type. +This happens primarily through the :py:class:`flytekit.extend.TypeEngine`. This engine works by invoking a series of :py:class:`TypeTransformers `. Each transformer is responsible for providing the functionality that the engine needs for a given native Python type. ***************** Callable Entities ***************** -Tasks, workflows, and launch plans form the core of the Flyte user experience. Each of these concepts are backed by one or more Python classes. These classes in turn, are instantiated by decorators (in the case of tasks and workflow) or a normal Python call (in the case of launch plans). +:ref:`Tasks `, :ref:`workflows `, and `launch plans ` form the core of the Flyte user experience. Each of these concepts are backed by one or more Python classes. These classes in turn, are instantiated by decorators (in the case of tasks and workflow) or a normal Python call (in the case of launch plans). Tasks ===== @@ -49,7 +49,7 @@ Please see the documentation on each of the classes for details. Workflows ========= -There are two workflow classes, which both inherit from the :py:class:`WorkflowBase ` class. +There are two workflow classes, and both of them inherit from the :py:class:`WorkflowBase ` class. .. autoclass:: flytekit.core.workflow.PythonFunctionWorkflow :noindex: @@ -60,7 +60,7 @@ There are two workflow classes, which both inherit from the :py:class:`WorkflowB Launch Plan =========== -There is also only one :py:class:`LaunchPlan ` class. +There is only one :py:class:`LaunchPlan ` class. .. autoclass:: flytekit.core.launch_plan.LaunchPlan :noindex: @@ -68,14 +68,14 @@ There is also only one :py:class:`LaunchPlan ` +You can execute all of these Flyte entities, which returns a :class:`~flytekit.remote.workflow_execution.FlyteWorkflowExecution` object. +For more information on Flyte entities, see the See the :ref:`remote flyte entities ` reference. .. code-block:: python @@ -89,7 +89,7 @@ You can also pass in ``wait=True`` to the :meth:`~flytekit.remote.remote.FlyteRe Syncing Remote State ******************** -Use the :meth:`~flytekit.remote.remote.FlyteRemote.sync` method to sync the entity object's state with the remote state +Use the :meth:`~flytekit.remote.remote.FlyteRemote.sync` method to sync the entity object's state with the remote state: .. code-block:: python @@ -100,6 +100,6 @@ Use the :meth:`~flytekit.remote.remote.FlyteRemote.sync` method to sync the enti Inspecting Execution Objects **************************** -At any time you can inspect the inputs, outputs, completion status, error status, and other aspects of a workflow +At any time, you can inspect the inputs, outputs, completion status, error status, and other aspects of a workflow execution object. See the :ref:`remote execution objects ` reference for a list of all the available attributes. diff --git a/docs/source/design/execution.rst b/docs/source/design/execution.rst index 2635a6c7c6..bf72f6ad3a 100644 --- a/docs/source/design/execution.rst +++ b/docs/source/design/execution.rst @@ -3,7 +3,7 @@ ####################### Execution Time Support ####################### -Most of the tasks that are written in flytekit will be Python functions decorated with ``@task`` which turns the body of the function into a Flyte Task, capable of being run independently, or included in any number of workflows. The interaction between flytekit and these tasks do not end once they have been serialized and registered onto the Flyte control plane however. When compiled, the command that will be executed when the Task is run is hardcoded into the task definition itself. +Most of the tasks that are written in Flytekit will be Python functions decorated with ``@task`` which turns the body of the function into a Flyte task, capable of being run independently, or included in any number of workflows. The interaction between Flytekit and these tasks do not end once they have been serialized and registered onto the Flyte control plane however. When compiled, the command that will be executed when the task is run is hardcoded into the task definition itself. In the basic ``@task`` decorated function scenario, the command to be run will be something containing ``pyflyte-execute``, which is one of the CLIs discussed in that section. @@ -11,10 +11,10 @@ That command, if you were to inspect a serialized task, might look something lik flytekit_venv pyflyte-execute --task-module app.workflows.failing_workflows --task-name divider --inputs {{.input}} --output-prefix {{.outputPrefix}} --raw-output-data-prefix {{.rawOutputDataPrefix}} -The point of running this script, or rather the reason for having any Flyte-related logic at execution time, is purely to codify and streamline the interaction between Flyte the platform, and the function body comprising user code. That is, the Flyte CLI is responsible for +The point of running this script, or rather the reason for having any Flyte-related logic at execution time, is purely to codify and streamline the interaction between Flyte the platform, and the function body comprising user code. The Flyte CLI is responsible for: -* I/O. The templated ``--inputs`` and ``--output-prefix`` arguments in the example command above will be filled in by the Flyte execution engine with S3 path (in the case of an AWS deployment). The ``pyflyte`` script will download the inputs to the right location in the container, and upload the results to the ``output-prefix`` location. +* I/O: The templated ``--inputs`` and ``--output-prefix`` arguments in the example command above will be filled in by the Flyte execution engine with S3 path (in the case of an AWS deployment). The ``pyflyte`` script will download the inputs to the right location in the container, and upload the results to the ``output-prefix`` location. * Ensure that raw output data prefix configuration option, which is again filled in by the Flyte engine, is respected so that ``FlyteFile``, ``FlyteDirectory``, and ``FlyteSchema`` objects offload their data to the correct place. -* Capture and handle error reporting. Exceptions thrown in the course of task execution are captured and uploaded to the Flyte control plane for display in the Console. -* Set up of helper utilities like the ``statsd`` handle, logging and logging levels, etc. +* Capture and handle error reporting: Exceptions thrown in the course of task execution are captured and uploaded to the Flyte control plane for display on the Console. +* Set up helper utilities like the ``statsd`` handle, logging and logging levels, etc. * Ensure configuration options about the Flyte backend, which are passed through by the Flyte engine, are properly loaded in Python memory. diff --git a/docs/source/design/index.rst b/docs/source/design/index.rst index 7da87efba6..7313aa3618 100644 --- a/docs/source/design/index.rst +++ b/docs/source/design/index.rst @@ -4,13 +4,13 @@ Overview ############################ -Flytekit is comprised of a handful of different logical components, each discusssed in greater detail in each link +Flytekit is comprised of a handful of different logical components, each discusssed in greater detail below: * :ref:`Models Files ` - These are almost Protobuf generated files. * :ref:`Authoring ` - This provides the core Flyte authoring experiences, allowing users to write tasks, workflows, and launch plans. * :ref:`Control Plane ` - The code here allows users to interact with the control plane through Python objects. * :ref:`Execution ` - A small shim layer basically that handles interaction with the Flyte ecosystem at execution time. -* :ref:`CLIs and Clients ` - Command line tools users may find themselves interacting with and the control plane client the CLIs call. +* :ref:`CLIs and Clients ` - Command line tools users may interact with, and the control plane client the CLIs call. .. toctree:: :maxdepth: 1 diff --git a/docs/source/design/models.rst b/docs/source/design/models.rst index 63be15098f..63fed55ea2 100644 --- a/docs/source/design/models.rst +++ b/docs/source/design/models.rst @@ -8,7 +8,7 @@ Model Files Description *********** This section deals with the files in `models `__ folder. -These files are nothing more than better formatted versions of the generated Python code from `Flyte IDL `__. In the future, we hope to be able to improve the Python code generator sufficiently to avoid this manual work. +These files are better formatted versions of the generated Python code from `Flyte IDL `__. In the future, we hope to be able to improve the Python code generator sufficiently to avoid this manual work. The __only__ reason these files exist is only because the Protobuf generated Python code doesn't work well with IDEs. It doesn't offer code completion, doesn't offer argument completion, docstrings, etc. @@ -37,4 +37,4 @@ Each Python object should have a ``to_flyte_idl`` and a ``from_flyte_idl`` funct ********* Testing ********* -Please add unit tests testing the conversion logic to the ``tests/flytekit/unit/models`` folder.s +Please add unit tests to the ``tests/flytekit/unit/models`` folder to test the conversion logic. diff --git a/docs/source/index.rst b/docs/source/index.rst index f83cbaa9b2..a33dc6b430 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -7,9 +7,9 @@ Flytekit Python Reference ************************* -This section of the documentation provides more detailed descriptions of the high-level design of ``flytekit`` and an -API reference for specific usage details of python functions, classes, and decorators that you import to specify tasks, -build workflows, and extend ``flytekit``. +This section of the documentation provides detailed descriptions of the high-level design of ``Flytekit`` and an +API reference for specific usage details of Python functions, classes, and decorators that you import to specify tasks, +build workflows, and extend ``Flytekit``. Installation ============ @@ -62,7 +62,7 @@ Expected output: |book| API Reference |hands-helping| Community -.. NOTE: the caption text is important for the sphinx theme to correctly render the nav header +.. NOTE: The caption text is important for the Sphinx theme to correctly render the nav header .. https://github.com/flyteorg/furo .. toctree:: :maxdepth: -1 diff --git a/flytekit/types/directory/__init__.py b/flytekit/types/directory/__init__.py index 97d8ab57ce..6aa193d7ce 100644 --- a/flytekit/types/directory/__init__.py +++ b/flytekit/types/directory/__init__.py @@ -7,6 +7,7 @@ .. autosummary:: :toctree: generated/ + :template: file_types.rst FlyteDirectory TensorboardLogs @@ -21,7 +22,7 @@ tensorboard = typing.TypeVar("tensorboard") TensorboardLogs = FlyteDirectory[tensorboard] """ - This type can be used to denote that the output is a folder that contains logs that can be loaded in tensorboard. - this is usually the SummaryWriter output in pytorch or Keras callbacks which record the history readable by - tensorboard + This type can be used to denote that the output is a folder that contains logs that can be loaded in TensorBoard. + This is usually the SummaryWriter output in PyTorch or Keras callbacks which record the history readable by + TensorBoard. """ diff --git a/flytekit/types/file/__init__.py b/flytekit/types/file/__init__.py index 81796fc49e..81ea1037bb 100644 --- a/flytekit/types/file/__init__.py +++ b/flytekit/types/file/__init__.py @@ -7,6 +7,7 @@ .. autosummary:: :toctree: generated/ + :template: file_types.rst FlyteFile HDF5EncodedFile @@ -27,71 +28,50 @@ # The following section provides some predefined aliases for commonly used FlyteFile formats. # This makes their usage extremely simple for the users. Please keep the list sorted. - hdf5 = typing.TypeVar("hdf5") +#: This can be used to denote that the returned file is of type hdf5 and can be received by other tasks that +#: accept an hdf5 format. This is usually useful for serializing Tensorflow models HDF5EncodedFile = FlyteFile[hdf5] -""" - This can be used to denote that the returned file is of type hdf5 and can be received by other tasks that - accept an hdf5 format. This is usually useful for serializing Tensorflow models -""" html = typing.TypeVar("html") +#: Can be used to receive or return an PNGImage. The underlying type is a FlyteFile type. This is just a +#: decoration and useful for attaching content type information with the file and automatically documenting code. HTMLPage = FlyteFile[html] -""" - Can be used to receive or return an PNGImage. The underlying type is a FlyteFile, type. This is just a - decoration and useful for attaching content type information with the file and automatically documenting code. -""" joblib = typing.TypeVar("joblib") +#: This File represents a file that was serialized using `joblib.dump` method can be loaded back using `joblib.load`. JoblibSerializedFile = FlyteFile[joblib] -""" - This File represents a file that was serialized using `joblib.dump` method can be loaded back using `joblib.load` -""" jpeg = typing.TypeVar("jpeg") +#: Can be used to receive or return an JPEGImage. The underlying type is a FlyteFile type. This is just a +#: decoration and useful for attaching content type information with the file and automatically documenting code. JPEGImageFile = FlyteFile[jpeg] -""" - Can be used to receive or return an JPEGImage. The underlying type is a FlyteFile, type. This is just a - decoration and useful for attaching content type information with the file and automatically documenting code. -""" pdf = typing.TypeVar("pdf") +#: Can be used to receive or return an PDFFile. The underlying type is a FlyteFile type. This is just a +#: decoration and useful for attaching content type information with the file and automatically documenting code. PDFFile = FlyteFile[pdf] -""" - Can be used to receive or return an PDFFile. The underlying type is a FlyteFile, type. This is just a - decoration and useful for attaching content type information with the file and automatically documenting code. -""" png = typing.TypeVar("png") +#: Can be used to receive or return an PNGImage. The underlying type is a FlyteFile type. This is just a +#: decoration and useful for attaching content type information with the file and automatically documenting code. PNGImageFile = FlyteFile[png] -""" - Can be used to receive or return an PNGImage. The underlying type is a FlyteFile, type. This is just a - decoration and useful for attaching content type information with the file and automatically documenting code. -""" python_pickle = typing.TypeVar("python_pickle") +#: This type can be used when a serialized Python pickled object is returned and shared between tasks. This only +#: adds metadata to the file in Flyte, but does not really carry any object information. PythonPickledFile = FlyteFile[python_pickle] -""" - This type can be used when a serialized python pickled object is returned and shared between tasks. This only - adds metadata to the file in Flyte, but does not really carry any object information -""" ipynb = typing.TypeVar("ipynb") +#: This type is used to identify a Python notebook file. PythonNotebook = FlyteFile[ipynb] -""" - This type is used to identify a python notebook file -""" svg = typing.TypeVar("svg") +#: Can be used to receive or return an SVGImage. The underlying type is a FlyteFile type. This is just a +#: decoration and useful for attaching content type information with the file and automatically documenting code. SVGImageFile = FlyteFile[svg] -""" - Can be used to receive or return an SVGImage. The underlying type is a FlyteFile, type. This is just a - decoration and useful for attaching content type information with the file and automatically documenting code. -""" csv = typing.TypeVar("csv") +#: Can be used to receive or return a CSVFile. The underlying type is a FlyteFile type. This is just a +#: decoration and useful for attaching content type information with the file and automatically documenting code. CSVFile = FlyteFile[csv] -""" - Can be used to receive or return a CSVFile. The underlying type is a FlyteFile, type. This is just a - decoration and useful for attaching content type information with the file and automatically documenting code. -""" diff --git a/flytekit/types/file/file.py b/flytekit/types/file/file.py index 8c5d4c6834..a215b6b851 100644 --- a/flytekit/types/file/file.py +++ b/flytekit/types/file/file.py @@ -31,10 +31,10 @@ class FlyteFile(os.PathLike, typing.Generic[T]): """ Since there is no native Python implementation of files and directories for the Flyte Blob type, (like how int exists for Flyte's Integer type) we need to create one so that users can express that their tasks take - in or return a file. There is ``pathlib.Path`` of course, (which is usable in flytekit as a return value, though + in or return a file. There is ``pathlib.Path`` of course, (which is usable in Flytekit as a return value, though not a return type), but it made more sense to create a new type esp. since we can add on additional properties. - Files (and directories) differ from the primitive types like floats and string in that flytekit typically uploads + Files (and directories) differ from the primitive types like floats and string in that Flytekit typically uploads the contents of the files to the blob store connected with your Flyte installation. That is, the Python native literal that represents a file is typically just the path to the file on the local filesystem. However in Flyte, an instance of a file is represented by a :py:class:`Blob ` literal, @@ -50,12 +50,12 @@ class FlyteFile(os.PathLike, typing.Generic[T]): In short, if a task returns ``"/path/to/file"`` and the task's signature is set to return ``FlyteFile``, then the contents of ``/path/to/file`` are uploaded. - You can also make it so that the upload does not happen. There are a few different types you use for + You can also make it so that the upload does not happen. There are different types of task/workflow signatures. Keep in mind that in the backend, in Admin and in the blob store, there is only one type that represents files, the :py:class:`Blob ` type. - Whether or not the uploading happens, and the behavior of the translation between Python native values and Flyte - literal values depends on a few things: + Whether the uploading happens or not, the behavior of the translation between Python native values and Flyte + literal values depends on a few attributes: * The declared Python type in the signature. These can be * :class:`python:flytekit.FlyteFile` @@ -65,9 +65,9 @@ class FlyteFile(os.PathLike, typing.Generic[T]): * :py:class:`flytekit.FlyteFile` * :py:class:`pathlib.Path` * :py:class:`str` - * Whether the value being converted is a "remote" path or not. For instance if a task returns a value of + * Whether the value being converted is a "remote" path or not. For instance, if a task returns a value of "http://www.google.com" as a ``FlyteFile``, obviously it doesn't make sense for us to try to upload that to the - Flyte blob store. So no remote paths are uploaded. flytekit considers a path remote if it starts with ``s3://``, + Flyte blob store. So no remote paths are uploaded. Flytekit considers a path remote if it starts with ``s3://``, ``gs://``, ``http(s)://``, or even ``file://``. ----------- @@ -208,7 +208,7 @@ def remote_path(self) -> os.PathLike: @property def remote_source(self) -> str: """ - If this is an input to a task, and the original path is ``s3://something``, flytekit will download the + If this is an input to a task, and the original path is an ``s3`` bucket, Flytekit downloads the file for the user. In case the user wants access to the original path, it will be here. """ return typing.cast(str, self._remote_source) diff --git a/flytekit/types/schema/types.py b/flytekit/types/schema/types.py index 0afed18158..71f1ef4ec7 100644 --- a/flytekit/types/schema/types.py +++ b/flytekit/types/schema/types.py @@ -24,7 +24,7 @@ class SchemaFormat(Enum): """ - Represents the the schema storage format (at rest). + Represents the schema storage format (at rest). Currently only parquet is supported """ @@ -264,7 +264,7 @@ def open( self, dataframe_fmt: type = pandas.DataFrame, override_mode: SchemaOpenMode = None ) -> typing.Union[SchemaReader, SchemaWriter]: """ - Will return a reader or writer depending on the mode of the object when created. This mode can be + Returns a reader or writer depending on the mode of the object when created. This mode can be overridden, but will depend on whether the override can be performed. For example, if the Object was created in a read-mode a "write mode" override is not allowed. if the object was created in write-mode, a read is allowed. From 35a4b9862e170e43c845de98df4cb57b352740df Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Wed, 23 Feb 2022 10:19:12 -0800 Subject: [PATCH 094/128] Bump idl (#862) Signed-off-by: Yee Hing Tong --- .github/workflows/pythonbuild.yml | 7 +---- dev-requirements.txt | 31 +++++++++---------- doc-requirements.txt | 24 +++++++------- requirements-spark2.txt | 22 ++++++------- requirements.txt | 26 ++++++++-------- setup.py | 1 - .../workflows/requirements.txt | 4 +-- 7 files changed, 51 insertions(+), 64 deletions(-) diff --git a/.github/workflows/pythonbuild.yml b/.github/workflows/pythonbuild.yml index 970673e090..5af4b5f634 100644 --- a/.github/workflows/pythonbuild.yml +++ b/.github/workflows/pythonbuild.yml @@ -138,12 +138,7 @@ jobs: - name: ShellCheck uses: ludeeus/action-shellcheck@master with: -<<<<<<< HEAD - ignore: boilerplate -======= - ignore_paths: - boilerplate ->>>>>>> 0d034c15 (Fixed format alias in Flytekit docs (#844)) + ignore_paths: boilerplate docs: runs-on: ubuntu-latest diff --git a/dev-requirements.txt b/dev-requirements.txt index 43568eb535..65e61b91fd 100644 --- a/dev-requirements.txt +++ b/dev-requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.9 +# This file is autogenerated by pip-compile with python 3.8 # To update, run: # # make dev-requirements.txt @@ -72,10 +72,7 @@ croniter==1.3.4 # -c requirements.txt # flytekit cryptography==36.0.1 - # via - # -c requirements.txt - # paramiko - # secretstorage + # via paramiko dataclasses-json==0.5.6 # via # -c requirements.txt @@ -116,6 +113,10 @@ docstring-parser==0.13 # flytekit filelock==3.6.0 # via virtualenv +flyteidl==0.22.3 + # via + # -c requirements.txt + # flytekit google-api-core[grpc]==2.5.0 # via # google-cloud-bigquery @@ -137,6 +138,8 @@ google-resumable-media==2.3.1 # via google-cloud-bigquery googleapis-common-protos==1.55.0 # via + # -c requirements.txt + # flyteidl # google-api-core # grpcio-status grpcio==1.44.0 @@ -160,11 +163,6 @@ importlib-metadata==4.11.2 # keyring iniconfig==1.1.1 # via pytest -jeepney==0.7.1 - # via - # -c requirements.txt - # keyring - # secretstorage jinja2==3.0.3 # via # -c requirements.txt @@ -256,6 +254,11 @@ protobuf==3.19.4 # googleapis-common-protos # grpcio-status # proto-plus + # protoc-gen-swagger +protoc-gen-swagger==0.1.0 + # via + # -c requirements.txt + # flyteidl py==1.11.0 # via # -c requirements.txt @@ -272,9 +275,7 @@ pyasn1==0.4.8 pyasn1-modules==0.2.8 # via google-auth pycparser==2.21 - # via - # -c requirements.txt - # cffi + # via cffi pynacl==1.5.0 # via paramiko pyparsing==3.0.7 @@ -348,10 +349,6 @@ retry==0.9.2 # flytekit rsa==4.8 # via google-auth -secretstorage==3.3.1 - # via - # -c requirements.txt - # keyring six==1.16.0 # via # -c requirements.txt diff --git a/doc-requirements.txt b/doc-requirements.txt index caee82a1ed..b8a2f19d65 100644 --- a/doc-requirements.txt +++ b/doc-requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.9 +# This file is autogenerated by pip-compile with python 3.8 # To update, run: # # make doc-requirements.txt @@ -42,9 +42,7 @@ cookiecutter==1.7.3 croniter==1.3.4 # via flytekit cryptography==36.0.1 - # via - # -r doc-requirements.in - # secretstorage + # via -r doc-requirements.in css-html-js-minify==2.5.5 # via sphinx-material dataclasses-json==0.5.6 @@ -63,8 +61,12 @@ docutils==0.17.1 # via # sphinx # sphinx-panels +flyteidl==0.22.3 + # via flytekit furo @ git+https://github.com/flyteorg/furo@main # via -r doc-requirements.in +googleapis-common-protos==1.54.0 + # via flyteidl grpcio==1.44.0 # via # -r doc-requirements.in @@ -77,10 +79,6 @@ importlib-metadata==4.11.2 # via # keyring # sphinx -jeepney==0.7.1 - # via - # keyring - # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -121,7 +119,13 @@ pandas==1.4.1 poyo==0.5.0 # via cookiecutter protobuf==3.19.4 - # via flytekit + # via + # flyteidl + # flytekit + # googleapis-common-protos + # protoc-gen-swagger +protoc-gen-swagger==0.1.0 + # via flyteidl py==1.11.0 # via retry pyarrow==6.0.1 @@ -167,8 +171,6 @@ responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit -secretstorage==3.3.1 - # via keyring six==1.16.0 # via # cookiecutter diff --git a/requirements-spark2.txt b/requirements-spark2.txt index 9ea67e42fc..3a007799e6 100644 --- a/requirements-spark2.txt +++ b/requirements-spark2.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.9 +# This file is autogenerated by pip-compile with python 3.8 # To update, run: # # make requirements-spark2.txt @@ -20,8 +20,6 @@ binaryornot==0.4.4 # via cookiecutter certifi==2021.10.8 # via requests -cffi==1.15.0 - # via cryptography chardet==4.0.0 # via binaryornot charset-normalizer==2.0.12 @@ -38,8 +36,6 @@ cookiecutter==1.7.3 # via flytekit croniter==1.3.4 # via flytekit -cryptography==36.0.1 - # via secretstorage dataclasses-json==0.5.6 # via flytekit decorator==5.1.1 @@ -52,16 +48,16 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit +flyteidl==0.22.3 + # via flytekit +googleapis-common-protos==1.54.0 + # via flyteidl grpcio==1.44.0 # via flytekit idna==3.3 # via requests importlib-metadata==4.11.2 # via keyring -jeepney==0.7.1 - # via - # keyring - # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -102,12 +98,14 @@ protobuf==3.19.4 # via # flyteidl # flytekit + # googleapis-common-protos + # protoc-gen-swagger +protoc-gen-swagger==0.1.0 + # via flyteidl py==1.11.0 # via retry pyarrow==6.0.1 # via flytekit -pycparser==2.21 - # via cffi pyrsistent==0.18.1 # via jsonschema python-dateutil==2.8.2 @@ -139,8 +137,6 @@ responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit -secretstorage==3.3.1 - # via keyring six==1.16.0 # via # cookiecutter diff --git a/requirements.txt b/requirements.txt index 8152af50b8..4c51a9362a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.9 +# This file is autogenerated by pip-compile with python 3.8 # To update, run: # # make requirements.txt @@ -18,8 +18,6 @@ binaryornot==0.4.4 # via cookiecutter certifi==2021.10.8 # via requests -cffi==1.15.0 - # via cryptography chardet==4.0.0 # via binaryornot charset-normalizer==2.0.12 @@ -36,8 +34,6 @@ cookiecutter==1.7.3 # via flytekit croniter==1.3.4 # via flytekit -cryptography==36.0.1 - # via secretstorage dataclasses-json==0.5.6 # via flytekit decorator==5.1.1 @@ -50,16 +46,16 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit +flyteidl==0.22.3 + # via flytekit +googleapis-common-protos==1.54.0 + # via flyteidl grpcio==1.44.0 # via flytekit idna==3.3 # via requests importlib-metadata==4.11.2 # via keyring -jeepney==0.7.1 - # via - # keyring - # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -100,12 +96,14 @@ protobuf==3.19.4 # via # flyteidl # flytekit + # googleapis-common-protos + # protoc-gen-swagger +protoc-gen-swagger==0.1.0 + # via flyteidl py==1.11.0 # via retry pyarrow==6.0.1 # via flytekit -pycparser==2.21 - # via cffi pyrsistent==0.18.1 # via jsonschema python-dateutil==2.8.2 @@ -116,7 +114,11 @@ python-dateutil==2.8.2 # pandas python-json-logger==2.0.2 # via flytekit +<<<<<<< HEAD python-slugify==6.1.1 +======= +python-slugify==6.1.0 +>>>>>>> 6a1ceea5 (Bump idl (#862)) # via cookiecutter pytimeparse==1.1.8 # via flytekit @@ -137,8 +139,6 @@ responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit -secretstorage==3.3.1 - # via keyring six==1.16.0 # via # cookiecutter diff --git a/setup.py b/setup.py index 7b08234cfd..438fd9a1a6 100644 --- a/setup.py +++ b/setup.py @@ -46,7 +46,6 @@ "python-dateutil>=2.1", "grpcio>=1.3.0,<2.0", "protobuf>=3.6.1,<4", - "protoc_gen_swagger", "python-json-logger>=2.0.0", "pytimeparse>=1.1.8,<2.0.0", "pytz", diff --git a/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt b/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt index fb365ed51d..22200b321a 100644 --- a/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt +++ b/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.9 +# This file is autogenerated by pip-compile with python 3.8 # To update, run: # # make tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt @@ -119,8 +119,6 @@ py==1.11.0 # via retry pyarrow==6.0.1 # via flytekit -pycparser==2.21 - # via cffi pyparsing==3.0.7 # via # matplotlib From c9800438dd1745764d34fbd2d359124fb68de142 Mon Sep 17 00:00:00 2001 From: SmritiSatyanV <94349093+SmritiSatyanV@users.noreply.github.com> Date: Fri, 25 Feb 2022 15:38:32 +0530 Subject: [PATCH 095/128] Updated authoring.rst (#863) * Updated authoring.rst Added directive Rephrased sentence * Fixed build error Signed-off-by: SmritiSatyanV * test-build-1 Signed-off-by: SmritiSatyanV Signed-off-by: maximsmol --- doc-requirements.txt | 7 +- docs/source/design/authoring.rst | 65 ++++++++++--------- .../integration/remote/test_remote.py | 2 +- .../flytekit/unit/extras/sqlite3/test_task.py | 4 +- 4 files changed, 41 insertions(+), 37 deletions(-) diff --git a/doc-requirements.txt b/doc-requirements.txt index b8a2f19d65..f734be7406 100644 --- a/doc-requirements.txt +++ b/doc-requirements.txt @@ -65,7 +65,7 @@ flyteidl==0.22.3 # via flytekit furo @ git+https://github.com/flyteorg/furo@main # via -r doc-requirements.in -googleapis-common-protos==1.54.0 +googleapis-common-protos==1.55.0 # via flyteidl grpcio==1.44.0 # via @@ -125,7 +125,10 @@ protobuf==3.19.4 # googleapis-common-protos # protoc-gen-swagger protoc-gen-swagger==0.1.0 - # via flyteidl + # via + # flyteidl + # flytekit + py==1.11.0 # via retry pyarrow==6.0.1 diff --git a/docs/source/design/authoring.rst b/docs/source/design/authoring.rst index aaa8f461cf..77b8b53628 100644 --- a/docs/source/design/authoring.rst +++ b/docs/source/design/authoring.rst @@ -1,33 +1,32 @@ .. _design-authoring: -############################ +####################### Authoring Structure -############################ +####################### -Enabling users to write tasks and workflows is the core feature of Flytekit, it is why it exists. This document goes over how some of the internals work. +One of the core features of Flytekit is to enable users to write tasks and workflows. In this section, we will understand how it works internally. -************* -Background -************* -Please refer `design doc `__. +.. note:: + + Please refer to the `design doc `__. ********************* Types and Type Engine ********************* -Flyte has its own type system, which is codified `in the IDL `__. Python, of course, has its own typing system, even though it's a dynamic language, and is mostly explained in `PEP 484 `_. In order to work properly, Flytekit needs to be able to convert between the two. +Flyte has its own type system, which is codified `in the IDL `__. Python has its own type system despite being a dynamic language, which is primarily explained in `PEP 484 `_. Flytekit needs to build a medium to bridge the gap between these two type systems. Type Engine ============= -This happens primarily through the :py:class:`flytekit.extend.TypeEngine`. This engine works by invoking a series of :py:class:`TypeTransformers `. Each transformer is responsible for providing the functionality that the engine needs for a given native Python type. +This primariliy happens through the :py:class:`flytekit.extend.TypeEngine`. This engine works by invoking a series of :py:class:`TypeTransformers `. Each transformer is responsible for providing the functionality that the engine requires for a given native Python type. ***************** Callable Entities ***************** -:ref:`Tasks `, :ref:`workflows `, and `launch plans ` form the core of the Flyte user experience. Each of these concepts are backed by one or more Python classes. These classes in turn, are instantiated by decorators (in the case of tasks and workflow) or a normal Python call (in the case of launch plans). +:ref:`Tasks `, :ref:`workflows `, and :ref:`launch plans ` form the core of the Flyte user experience. Each of these concepts is backed by one or more Python classes. These classes in turn, are instantiated by decorators (in the case of tasks and workflow) or a regular Python call (in the case of launch plans). Tasks ===== -This is the current task class hierarchy. +This is the current task class hierarchy: .. inheritance-diagram:: flytekit.core.python_function_task.PythonFunctionTask flytekit.core.python_function_task.PythonInstanceTask flytekit.extras.sqlite3.task.SQLite3Task :parts: 1 @@ -47,9 +46,10 @@ Please see the documentation on each of the classes for details. .. autoclass:: flytekit.core.python_function_task.PythonFunctionTask :noindex: + Workflows -========= -There are two workflow classes, and both of them inherit from the :py:class:`WorkflowBase ` class. +========== +There are two workflow classes, and both inherit from the :py:class:`WorkflowBase ` class. .. autoclass:: flytekit.core.workflow.PythonFunctionWorkflow :noindex: @@ -65,17 +65,17 @@ There is only one :py:class:`LaunchPlan ` .. autoclass:: flytekit.core.launch_plan.LaunchPlan :noindex: +.. exception_handling: + ****************** Exception Handling ****************** -Exception handling is done along two dimensions: +Exception handling takes place along two dimensions: -* System vs User: We try to differentiate between user exceptions and Flytekit/system level exceptions. For instance, if Flytekit - fails to upload its outputs, that's a system exception. If you the user raise a ``ValueError`` because of unexpected input - in the task code, that's a user exception. -* Recoverable vs Non-recoverable: Recoverable errors are retried and counted against your task's retry count. Non-recoverable errors just fail. System exceptions, by default, are recoverable (since there's a good chance it was just a blip). +* System vs. User: We try to differentiate between user exceptions and Flytekit/system-level exceptions. For instance, if Flytekit fails to upload its outputs, that's a system exception. If the user raises a ``ValueError`` because of an unexpected input in the task code, that's a user exception. +* Recoverable vs. Non-recoverable: Recoverable errors will be retried and counted against the task's retry count. Non-recoverable errors will simply fail. System exceptions are by default recoverable (since there's a good chance it was just a blip). -This is the user exception tree. Feel free to raise any of these exception classes. Note that the ``FlyteRecoverableException`` is the only recoverable one. All others, along with all non-Flytekit defined exceptions, are non-recoverable. +Here's the user exception tree. Feel free to raise any of these exception classes. Note that the ``FlyteRecoverableException`` is the only recoverable exception. All others, along with all the non-Flytekit defined exceptions, are non-recoverable. .. inheritance-diagram:: flytekit.common.exceptions.user.FlyteValidationException flytekit.common.exceptions.user.FlyteEntityAlreadyExistsException flytekit.common.exceptions.user.FlyteValueException flytekit.common.exceptions.user.FlyteTimeout flytekit.common.exceptions.user.FlyteAuthenticationException flytekit.common.exceptions.user.FlyteRecoverableException :parts: 1 @@ -83,8 +83,8 @@ This is the user exception tree. Feel free to raise any of these exception class Implementation ============== -For those who want to dig a bit deeper, take a look at the :py:class:`flytekit.common.exceptions.scopes.FlyteScopedException` classes. -There are two decorators which you'll find interspersed throughout the codebase. +For those who want to dig deeper, take a look at the :py:class:`flytekit.common.exceptions.scopes.FlyteScopedException` classes. +There are two decorators that are interspersed throughout the codebase. .. autofunction:: flytekit.common.exceptions.scopes.system_entry_point @@ -93,18 +93,19 @@ There are two decorators which you'll find interspersed throughout the codebase. ************** Call Patterns ************** -The three entities above are all callable. In Flyte terms, this means they can be invoked to yield a unit (or units) of work. -In Python terms, this means you can add ``()`` to the end of one of it which invokes the ``__call__`` method on the object. +The above-mentioned entities (tasks, workflows, and launch plan) are callable. They can be invoked to yield a unit (or units) of work in Flyte. + +In Pythonic terms, when you add ``()`` to the end of one of the entities, it invokes the ``__call__`` method on the object. What happens when a callable entity is called depends on the current context, specifically the current :py:class:`flytekit.FlyteContext` Raw Task Execution -================== -This is what happens when a task is just run as part of a unit test. The ``@task`` decorator actually turns the decorated function into an instance of the ``PythonFunctionTask`` object but when a user calls it, ``task1()``, outside of a workflow, the original function is called without interference by flytekit. +=================== +This is what happens when a task is just run as part of a unit test. The ``@task`` decorator actually turns the decorated function into an instance of the ``PythonFunctionTask`` object, but when a user calls the ``task()`` outside of a workflow, the original function is called without any interference by Flytekit. Task Execution Inside Workflow -============================== -This is what happens, *to the task* when a workflow is being run locally, say as part of a unit test for the workflow. +=============================== +When a workflow is run locally (say as a part of a unit test), certain changes occur in the ``task``. Before going further, there is a special object that's worth mentioning, the :py:class:`flytekit.extend.Promise`. @@ -127,19 +128,19 @@ Let's assume we have a workflow like :: d = t2(a=y, b=b) return x, d -As discussed in the Promise object's documentation, when a task is called from inside a workflow, the Python native values that the raw underlying functions return are first converted into Flyte IDL literals, and then wrapped inside ``Promise`` objects. One ``Promise`` is created for every return variable. +As discussed in the Promise object's documentation, when a task is called from inside a workflow, the Python native values returned by the raw underlying functions are first converted into Flyte IDL literals and then wrapped inside ``Promise`` objects. One ``Promise`` is created for every return variable. -When the next task is called, logic is triggered to unwrap these promises. +When the next task is called, the logic is triggered to unwrap these Promises. Compilation =========== -When a workflow is compiled, instead of producing ``Promise`` objects that wrap literal values, they wrap a :py:class:`flytekit.core.promise.NodeOutput` instead. This is how data dependency is tracked between tasks. +When a workflow is compiled, instead of producing Promise objects that wrap literal values, they wrap a :py:class:`flytekit.core.promise.NodeOutput` instead. This helps track data dependency between tasks. Branch Skip =========== -If it's been determined that a conditional is not true, then Flytekit will skip calling the task. This way, any side-effects in the task logic will not be run. +If a :py:func:`flytekit.conditional` is determined to be false, then Flytekit will skip calling the task. This avoids running the unintended task. .. note:: - Even though in the discussion above, we talked about a task's execution pattern, the same actually applied to workflows and launch plans. + We discussed about a task's execution pattern above. The same pattern can be applied to workflows and launch plans too! diff --git a/tests/flytekit/integration/remote/test_remote.py b/tests/flytekit/integration/remote/test_remote.py index ba105d386b..0ca4f26b1e 100644 --- a/tests/flytekit/integration/remote/test_remote.py +++ b/tests/flytekit/integration/remote/test_remote.py @@ -260,7 +260,7 @@ def test_execute_python_workflow_list_of_floats(flyteclient, flyte_workflows_reg def test_execute_sqlite3_task(flyteclient, flyte_workflows_register, flyte_remote_env): remote = FlyteRemote.from_config(PROJECT, "development") - example_db = "https://cdn.sqlitetutorial.net/wp-content/uploads/2018/03/chinook.zip" + example_db = "https://www.sqlitetutorial.net/wp-content/uploads/2018/03/chinook.zip" interactive_sql_task = SQLite3Task( "basic_querying", query_template="select TrackId, Name from tracks limit {{.inputs.limit}}", diff --git a/tests/flytekit/unit/extras/sqlite3/test_task.py b/tests/flytekit/unit/extras/sqlite3/test_task.py index 365a55cee8..f586d94a16 100644 --- a/tests/flytekit/unit/extras/sqlite3/test_task.py +++ b/tests/flytekit/unit/extras/sqlite3/test_task.py @@ -6,7 +6,7 @@ # https://www.sqlitetutorial.net/sqlite-sample-database/ from flytekit.types.schema import FlyteSchema -EXAMPLE_DB = "https://cdn.sqlitetutorial.net/wp-content/uploads/2018/03/chinook.zip" +EXAMPLE_DB = "https://www.sqlitetutorial.net/wp-content/uploads/2018/03/chinook.zip" # This task belongs to test_task_static but is intentionally here to help test tracking tk = SQLite3Task( @@ -28,7 +28,7 @@ def test_task_static(): def test_task_schema(): # sqlite3_start - DB_LOCATION = "https://cdn.sqlitetutorial.net/wp-content/uploads/2018/03/chinook.zip" + DB_LOCATION = "https://www.sqlitetutorial.net/wp-content/uploads/2018/03/chinook.zip" sql_task = SQLite3Task( "test", From 517de983260b287b0a04c7ffb2aec77e2e8d2329 Mon Sep 17 00:00:00 2001 From: SmritiSatyanV <94349093+SmritiSatyanV@users.noreply.github.com> Date: Fri, 25 Feb 2022 17:07:48 +0530 Subject: [PATCH 096/128] Updated authoring.rst (#866) Added the directive correctly Signed-off-by: SmritiSatyanV Signed-off-by: maximsmol --- docs/source/design/authoring.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/design/authoring.rst b/docs/source/design/authoring.rst index 77b8b53628..ee676cf165 100644 --- a/docs/source/design/authoring.rst +++ b/docs/source/design/authoring.rst @@ -65,7 +65,7 @@ There is only one :py:class:`LaunchPlan ` .. autoclass:: flytekit.core.launch_plan.LaunchPlan :noindex: -.. exception_handling: +.. _exception_handling: ****************** Exception Handling From e01e7c0182c5d81d6c5062abbe7d9a29351870c2 Mon Sep 17 00:00:00 2001 From: Matthew Griffin <1matthewgriffin@gmail.com> Date: Sat, 26 Feb 2022 10:44:40 -0600 Subject: [PATCH 097/128] Change docs for HTMLPage type to say HTMLPage instead of PNGImage (#868) Signed-off-by: maximsmol --- flytekit/types/file/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flytekit/types/file/__init__.py b/flytekit/types/file/__init__.py index 81ea1037bb..44841d7e35 100644 --- a/flytekit/types/file/__init__.py +++ b/flytekit/types/file/__init__.py @@ -34,7 +34,7 @@ HDF5EncodedFile = FlyteFile[hdf5] html = typing.TypeVar("html") -#: Can be used to receive or return an PNGImage. The underlying type is a FlyteFile type. This is just a +#: Can be used to receive or return an HTMLPage. The underlying type is a FlyteFile type. This is just a #: decoration and useful for attaching content type information with the file and automatically documenting code. HTMLPage = FlyteFile[html] From 75e4da9a37e4a48fb208d3b5ed02c688987fb773 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Wed, 2 Mar 2022 05:14:14 +0800 Subject: [PATCH 098/128] Revisit StructuredDatasetDecoder interface (#865) Signed-off-by: Kevin Su Signed-off-by: maximsmol --- flytekit/models/literals.py | 6 +-- flytekit/types/structured/basic_dfs.py | 14 +++--- flytekit/types/structured/bigquery.py | 19 ++++---- .../types/structured/structured_dataset.py | 35 ++++++--------- .../flytekitplugins/fsspec/arrow.py | 7 ++- .../flytekitplugins/fsspec/pandas.py | 7 ++- .../flytekitplugins/spark/sd_transformers.py | 4 ++ plugins/flytekit-spark/tests/test_wf.py | 43 +++++++++---------- .../unit/core/test_structured_dataset.py | 6 ++- .../core/test_structured_dataset_handlers.py | 10 ++--- .../test_structured_dataset_workflow.py | 2 + 11 files changed, 75 insertions(+), 78 deletions(-) diff --git a/flytekit/models/literals.py b/flytekit/models/literals.py index 6584334450..c4ac6ca289 100644 --- a/flytekit/models/literals.py +++ b/flytekit/models/literals.py @@ -1,5 +1,5 @@ -import typing from datetime import datetime as _datetime +from typing import Optional import pytz as _pytz from flyteidl.core import literals_pb2 as _literals_pb2 @@ -592,7 +592,7 @@ def from_flyte_idl(cls, pb2_object): class StructuredDatasetMetadata(_common.FlyteIdlEntity): - def __init__(self, structured_dataset_type: StructuredDatasetType = None): + def __init__(self, structured_dataset_type: Optional[StructuredDatasetType] = None): self._structured_dataset_type = structured_dataset_type @property @@ -614,7 +614,7 @@ def from_flyte_idl(cls, pb2_object: _literals_pb2.StructuredDatasetMetadata) -> class StructuredDataset(_common.FlyteIdlEntity): - def __init__(self, uri: str, metadata: typing.Optional[StructuredDatasetMetadata] = None): + def __init__(self, uri: str, metadata: Optional[StructuredDatasetMetadata] = None): """ A strongly typed schema that defines the interface of data retrieved from the underlying storage medium. """ diff --git a/flytekit/types/structured/basic_dfs.py b/flytekit/types/structured/basic_dfs.py index 49b2f13ed9..ff9d692cec 100644 --- a/flytekit/types/structured/basic_dfs.py +++ b/flytekit/types/structured/basic_dfs.py @@ -55,14 +55,13 @@ def decode( self, ctx: FlyteContext, flyte_value: literals.StructuredDataset, + current_task_metadata: StructuredDatasetMetadata, ) -> pd.DataFrame: path = flyte_value.uri local_dir = ctx.file_access.get_random_local_directory() ctx.file_access.get_data(path, local_dir, is_multipart=True) - if flyte_value.metadata.structured_dataset_type.columns: - columns = [] - for c in flyte_value.metadata.structured_dataset_type.columns: - columns.append(c.name) + if current_task_metadata.structured_dataset_type and current_task_metadata.structured_dataset_type.columns: + columns = [c.name for c in current_task_metadata.structured_dataset_type.columns] return pd.read_parquet(local_dir, columns=columns) return pd.read_parquet(local_dir) @@ -94,14 +93,13 @@ def decode( self, ctx: FlyteContext, flyte_value: literals.StructuredDataset, + current_task_metadata: StructuredDatasetMetadata, ) -> pa.Table: path = flyte_value.uri local_dir = ctx.file_access.get_random_local_directory() ctx.file_access.get_data(path, local_dir, is_multipart=True) - if flyte_value.metadata.structured_dataset_type.columns: - columns = [] - for c in flyte_value.metadata.structured_dataset_type.columns: - columns.append(c.name) + if current_task_metadata.structured_dataset_type and current_task_metadata.structured_dataset_type.columns: + columns = [c.name for c in current_task_metadata.structured_dataset_type.columns] return pq.read_table(local_dir, columns=columns) return pq.read_table(local_dir) diff --git a/flytekit/types/structured/bigquery.py b/flytekit/types/structured/bigquery.py index 923ea06e9e..aa0ef42f6b 100644 --- a/flytekit/types/structured/bigquery.py +++ b/flytekit/types/structured/bigquery.py @@ -11,7 +11,6 @@ from flytekit.models.types import StructuredDatasetType from flytekit.types.structured.structured_dataset import ( BIGQUERY, - DF, StructuredDataset, StructuredDatasetDecoder, StructuredDatasetEncoder, @@ -29,7 +28,9 @@ def _write_to_bq(structured_dataset: StructuredDataset): client.load_table_from_dataframe(df, table_id) -def _read_from_bq(flyte_value: literals.StructuredDataset) -> pd.DataFrame: +def _read_from_bq( + flyte_value: literals.StructuredDataset, current_task_metadata: StructuredDatasetMetadata +) -> pd.DataFrame: path = flyte_value.uri _, project_id, dataset_id, table_id = re.split("\\.|://|:", path) client = bigquery_storage.BigQueryReadClient() @@ -37,10 +38,8 @@ def _read_from_bq(flyte_value: literals.StructuredDataset) -> pd.DataFrame: parent = "projects/{}".format(project_id) read_options = None - if flyte_value.metadata.structured_dataset_type.columns: - columns = [] - for c in flyte_value.metadata.structured_dataset_type.columns: - columns.append(c.name) + if current_task_metadata.structured_dataset_type and current_task_metadata.structured_dataset_type.columns: + columns = [c.name for c in current_task_metadata.structured_dataset_type.columns] read_options = types.ReadSession.TableReadOptions(selected_fields=columns) requested_session = types.ReadSession(table=table, data_format=types.DataFormat.ARROW, read_options=read_options) @@ -78,8 +77,9 @@ def decode( self, ctx: FlyteContext, flyte_value: literals.StructuredDataset, - ) -> typing.Union[DF, typing.Generator[DF, None, None]]: - return _read_from_bq(flyte_value) + current_task_metadata: StructuredDatasetMetadata, + ) -> pd.DataFrame: + return _read_from_bq(flyte_value, current_task_metadata) class ArrowToBQEncodingHandlers(StructuredDatasetEncoder): @@ -106,7 +106,8 @@ def decode( self, ctx: FlyteContext, flyte_value: literals.StructuredDataset, - ) -> typing.Union[DF, typing.Generator[DF, None, None]]: + current_task_metadata: StructuredDatasetMetadata, + ) -> pa.Table: return pa.Table.from_pandas(_read_from_bq(flyte_value)) diff --git a/flytekit/types/structured/structured_dataset.py b/flytekit/types/structured/structured_dataset.py index 819bc012cc..fa57d7b5de 100644 --- a/flytekit/types/structured/structured_dataset.py +++ b/flytekit/types/structured/structured_dataset.py @@ -105,9 +105,7 @@ def all(self) -> DF: if self._dataframe_type is None: raise ValueError("No dataframe type set. Use open() to set the local dataframe type you want to use.") ctx = FlyteContextManager.current_context() - return flyte_dataset_transformer.open_as( - ctx, self.literal, self._dataframe_type, updated_metadata=self.metadata - ) + return flyte_dataset_transformer.open_as(ctx, self.literal, self._dataframe_type, self.metadata) def iter(self) -> Generator[DF, None, None]: if self._dataframe_type is None: @@ -261,14 +259,17 @@ def decode( self, ctx: FlyteContext, flyte_value: literals.StructuredDataset, + current_task_metadata: StructuredDatasetMetadata, ) -> Union[DF, Generator[DF, None, None]]: """ This is code that will be called by the dataset transformer engine to ultimately translate from a Flyte Literal value into a Python instance. - :param ctx: + :param ctx: A FlyteContext, useful in accessing the filesystem and other attributes :param flyte_value: This will be a Flyte IDL StructuredDataset Literal - do not confuse this with the StructuredDataset class defined also in this module. + :param current_task_metadata: Metadata object containing the type (and columns if any) for the currently + executing task. This type may have more or less information than the type information bundled inside the incoming flyte_value. :return: This function can either return an instance of the dataframe that this decoder handles, or an iterator of those dataframes. """ @@ -585,7 +586,7 @@ def t2(in_a: Annotated[StructuredDataset, kwtypes(col_b=float)]): ... sd._literal_sd = sd_literal return sd else: - return self.open_as(ctx, sd_literal, df_type=expected_python_type) + return self.open_as(ctx, sd_literal, expected_python_type, metad) # Start handling for StructuredDataset scalars, first look at the columns incoming_columns = lv.scalar.structured_dataset.metadata.structured_dataset_type.columns @@ -596,8 +597,7 @@ def t2(in_a: Annotated[StructuredDataset, kwtypes(col_b=float)]): ... if column_dict is None or len(column_dict) == 0: # but if it does, then we just copy it over if incoming_columns is not None and incoming_columns != []: - for c in incoming_columns: - final_dataset_columns.append(c) + final_dataset_columns = incoming_columns.copy() # If the current running task's input does have columns defined else: final_dataset_columns = self._convert_ordered_dict_of_columns_to_list(column_dict) @@ -631,22 +631,18 @@ def open_as( ctx: FlyteContext, sd: literals.StructuredDataset, df_type: Type[DF], - updated_metadata: Optional[StructuredDatasetMetadata] = None, + updated_metadata: StructuredDatasetMetadata, ) -> DF: """ - - :param ctx: + :param ctx: A FlyteContext, useful in accessing the filesystem and other attributes :param sd: :param df_type: - :param meta: New metadata type, since it might be different from the metadata in the literal. - :return: + :param updated_metadata: New metadata type, since it might be different from the metadata in the literal. + :return: dataframe. It could be pandas dataframe or arrow table, etc. """ protocol = protocol_prefix(sd.uri) decoder = self.get_decoder(df_type, protocol, sd.metadata.structured_dataset_type.format) - # todo: revisit this, we probably should add a new field to the decoder interface - if updated_metadata: - sd._metadata = updated_metadata - result = decoder.decode(ctx, sd) + result = decoder.decode(ctx, sd, updated_metadata) if isinstance(result, types.GeneratorType): raise ValueError(f"Decoder {decoder} returned iterator {result} but whole value requested from {sd}") return result @@ -656,14 +652,11 @@ def iter_as( ctx: FlyteContext, sd: literals.StructuredDataset, df_type: Type[DF], - updated_metadata: Optional[StructuredDatasetMetadata] = None, + updated_metadata: StructuredDatasetMetadata, ) -> Generator[DF, None, None]: protocol = protocol_prefix(sd.uri) decoder = self.DECODERS[df_type][protocol][sd.metadata.structured_dataset_type.format] - # todo: revisit this, should we add a new field to the decoder interface - if updated_metadata: - sd._metadata = updated_metadata - result = decoder.decode(ctx, sd) + result = decoder.decode(ctx, sd, updated_metadata) if not isinstance(result, types.GeneratorType): raise ValueError(f"Decoder {decoder} didn't return iterator {result} but should have from {sd}") return result diff --git a/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/arrow.py b/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/arrow.py index c17e2fc8bd..d47318666f 100644 --- a/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/arrow.py +++ b/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/arrow.py @@ -47,6 +47,7 @@ def decode( self, ctx: FlyteContext, flyte_value: literals.StructuredDataset, + current_task_metadata: StructuredDatasetMetadata, ) -> pa.Table: uri = flyte_value.uri if not ctx.file_access.is_remote(uri): @@ -54,10 +55,8 @@ def decode( _, path = split_protocol(uri) columns = None - if flyte_value.metadata.structured_dataset_type.columns: - columns = [] - for c in flyte_value.metadata.structured_dataset_type.columns: - columns.append(c.name) + if current_task_metadata.structured_dataset_type and current_task_metadata.structured_dataset_type.columns: + columns = [c.name for c in current_task_metadata.structured_dataset_type.columns] try: fs = FSSpecPersistence.get_filesystem(uri) return pq.read_table(path, filesystem=fs, columns=columns) diff --git a/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/pandas.py b/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/pandas.py index 52bcc4522a..07b58d243a 100644 --- a/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/pandas.py +++ b/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/pandas.py @@ -58,14 +58,13 @@ def decode( self, ctx: FlyteContext, flyte_value: literals.StructuredDataset, + current_task_metadata: StructuredDatasetMetadata, ) -> pd.DataFrame: uri = flyte_value.uri columns = None kwargs = get_storage_options(uri) - if flyte_value.metadata.structured_dataset_type.columns: - columns = [] - for c in flyte_value.metadata.structured_dataset_type.columns: - columns.append(c.name) + if current_task_metadata.structured_dataset_type and current_task_metadata.structured_dataset_type.columns: + columns = [c.name for c in current_task_metadata.structured_dataset_type.columns] try: return pd.read_parquet(uri, columns=columns, storage_options=kwargs) except NoCredentialsError: diff --git a/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py b/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py index 2466d3fc13..cd451fa080 100644 --- a/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py +++ b/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py @@ -39,8 +39,12 @@ def decode( self, ctx: FlyteContext, flyte_value: literals.StructuredDataset, + current_task_metadata: StructuredDatasetMetadata, ) -> DataFrame: user_ctx = FlyteContext.current_context().user_space_params + if current_task_metadata.structured_dataset_type and current_task_metadata.structured_dataset_type.columns: + columns = [c.name for c in current_task_metadata.structured_dataset_type.columns] + return user_ctx.spark_session.read.parquet(flyte_value.uri).select(*columns) return user_ctx.spark_session.read.parquet(flyte_value.uri) diff --git a/plugins/flytekit-spark/tests/test_wf.py b/plugins/flytekit-spark/tests/test_wf.py index a0a624fec7..8c42a6162f 100644 --- a/plugins/flytekit-spark/tests/test_wf.py +++ b/plugins/flytekit-spark/tests/test_wf.py @@ -6,6 +6,11 @@ from flytekit import kwtypes, task, workflow from flytekit.types.schema import FlyteSchema +try: + from typing import Annotated +except ImportError: + from typing_extensions import Annotated + def test_wf1_with_spark(): @task(task_config=Spark()) @@ -53,27 +58,6 @@ def my_wf() -> my_schema: assert df2 is not None -def test_ddwf1_with_spark(): - @task(task_config=Spark()) - def my_spark(a: int) -> (int, str): - session = flytekit.current_context().spark_session - assert session.sparkContext.appName == "FlyteSpark: ex:local:local:local" - return a + 2, "world" - - @task - def t2(a: str, b: str) -> str: - return b + a - - @workflow - def my_wf(a: int, b: str) -> (int, str): - x, y = my_spark(a=a) - d = t2(a=y, b=b) - return x, d - - x = my_wf(a=5, b="hello ") - assert x == (7, "hello world") - - def test_fs_sd_compatibility(): my_schema = FlyteSchema[kwtypes(name=str, age=int)] @@ -108,7 +92,6 @@ def test_spark_dataframe_return(): def my_spark(a: int) -> my_schema: session = flytekit.current_context().spark_session df = session.createDataFrame([("Alice", a)], my_schema.column_names()) - print(type(df)) return df @workflow @@ -120,3 +103,19 @@ def my_wf(a: int) -> my_schema: df2 = reader.all() result_df = df2.reset_index(drop=True) == pd.DataFrame(data={"name": ["Alice"], "age": [5]}).reset_index(drop=True) assert result_df.all().all() + + +def test_read_spark_subset_columns(): + @task + def t1() -> pd.DataFrame: + return pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]}) + + @task(task_config=Spark()) + def t2(df: Annotated[pyspark.sql.DataFrame, kwtypes(Name=str)]) -> int: + return len(df.columns) + + @workflow() + def wf() -> int: + return t2(df=t1()) + + assert wf() == 1 diff --git a/tests/flytekit/unit/core/test_structured_dataset.py b/tests/flytekit/unit/core/test_structured_dataset.py index a7ef1ea953..8efaecfcc9 100644 --- a/tests/flytekit/unit/core/test_structured_dataset.py +++ b/tests/flytekit/unit/core/test_structured_dataset.py @@ -221,6 +221,7 @@ def decode( self, ctx: FlyteContext, flyte_value: literals.StructuredDataset, + current_task_metadata: StructuredDatasetMetadata, ) -> typing.Union[typing.Generator[pd.DataFrame, None, None]]: yield pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]}) @@ -241,8 +242,9 @@ def decode( self, ctx: FlyteContext, flyte_value: literals.StructuredDataset, + current_task_metadata: StructuredDatasetMetadata, ) -> pd.DataFrame: - pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]}) + return pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]}) StructuredDatasetTransformerEngine.register( MockPandasDecodingHandlers(pd.DataFrame, "tmpfs"), default_for_type=False, override=True @@ -288,7 +290,7 @@ def test_to_python_value_with_incoming_columns(): # check when columns are not specified, should pull both and add column information. sd = fdt.to_python_value(ctx, lit, StructuredDataset) - assert sd.metadata.structured_dataset_type.columns[0].name == "age" + assert len(sd.metadata.structured_dataset_type.columns) == 2 # should also work if subset type is just an annotated pd.DataFrame subset_pd_type = Annotated[pd.DataFrame, kwtypes(age=int)] diff --git a/tests/flytekit/unit/core/test_structured_dataset_handlers.py b/tests/flytekit/unit/core/test_structured_dataset_handlers.py index 0d755ab78b..ada7483a0f 100644 --- a/tests/flytekit/unit/core/test_structured_dataset_handlers.py +++ b/tests/flytekit/unit/core/test_structured_dataset_handlers.py @@ -6,6 +6,7 @@ from flytekit.core import context_manager from flytekit.core.base_task import kwtypes +from flytekit.models.literals import StructuredDatasetMetadata from flytekit.models.types import StructuredDatasetType from flytekit.types.structured import basic_dfs from flytekit.types.structured.structured_dataset import ( @@ -26,12 +27,11 @@ def test_pandas(): decoder = basic_dfs.ParquetToPandasDecodingHandler("/") ctx = context_manager.FlyteContextManager.current_context() - sd = StructuredDataset( - dataframe=df, - ) - sd_lit = encoder.encode(ctx, sd, StructuredDatasetType(format="parquet")) + sd = StructuredDataset(dataframe=df) + sd_type = StructuredDatasetType(format="parquet") + sd_lit = encoder.encode(ctx, sd, sd_type) - df2 = decoder.decode(ctx, sd_lit) + df2 = decoder.decode(ctx, sd_lit, StructuredDatasetMetadata(sd_type)) assert df.equals(df2) diff --git a/tests/flytekit/unit/types/structured_dataset/test_structured_dataset_workflow.py b/tests/flytekit/unit/types/structured_dataset/test_structured_dataset_workflow.py index d911f971e4..0fc0bd976c 100644 --- a/tests/flytekit/unit/types/structured_dataset/test_structured_dataset_workflow.py +++ b/tests/flytekit/unit/types/structured_dataset/test_structured_dataset_workflow.py @@ -54,6 +54,7 @@ def decode( self, ctx: FlyteContext, flyte_value: literals.StructuredDataset, + current_task_metadata: StructuredDatasetMetadata, ) -> pd.DataFrame: return pd_df @@ -86,6 +87,7 @@ def decode( self, ctx: FlyteContext, flyte_value: literals.StructuredDataset, + current_task_metadata: StructuredDatasetMetadata, ) -> typing.Union[DF, typing.Generator[DF, None, None]]: path = flyte_value.uri local_dir = ctx.file_access.get_random_local_directory() From 35867aaed51e5a401ebf2aec1b90ba8578acb769 Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Tue, 1 Mar 2022 13:57:46 -0800 Subject: [PATCH 099/128] Remove legacy mentions in contributing guide (#870) Signed-off-by: Yee Hing Tong Signed-off-by: maximsmol --- docs/source/contributing.rst | 19 ++----------------- 1 file changed, 2 insertions(+), 17 deletions(-) diff --git a/docs/source/contributing.rst b/docs/source/contributing.rst index 1a69f67c8a..9b147de866 100644 --- a/docs/source/contributing.rst +++ b/docs/source/contributing.rst @@ -9,14 +9,10 @@ First off, thank you for thinking about contributing! Below you'll find instruct Please also take some time to read through the :std:ref:`design guides `, which describe the various parts of Flytekit and should make contributing easier. ******************* -📜 Quick Background +📜 Background ******************* -The first version of the Flytekit library was written circa 2017, before mypy typing was mainstream, and -targeted Python 2. That legacy code will be fully deprecated and removed in 2022. Since there are still -users of Flytekit that rely on that legacy API, you'll see 2 separate and distinct code paths within this repo. -Users and contributors should ignore the legacy sections. Below is a listing of the most important packages that -comprise the new API: +Below is a listing of the most important packages that comprise the flytekit SDK: - ``flytekit/core`` This holds all the core functionality of the new API. @@ -45,17 +41,6 @@ comprise the new API: makes compiled Flyte tasks and workflows less portable (i.e. if you run ``pyflyte package`` can someone else use those serialized objects). -Most of the other folders are for legacy Flytekit, support for which will be dropped in early 2022. For the most part, -please ignore the following folders: - -- ``flytekit/plugins`` -- ``flytekit/common`` - (the ``translator.py`` file is an exception) -- ``flytekit/engines`` -- ``flytekit/interfaces`` -- ``flytekit/sdk`` -- ``flytekit/type_engines`` - Please also see the :std:ref:`design overview section ` for more in-depth information. From 8636f36f5021d79c3551df95542c8103f82e5781 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Wed, 2 Mar 2022 20:18:10 +0800 Subject: [PATCH 100/128] Add GCS protocol in the structured dataset (#869) Signed-off-by: Kevin Su Signed-off-by: maximsmol --- flytekit/types/structured/basic_dfs.py | 3 ++- flytekit/types/structured/structured_dataset.py | 1 + .../flytekit-data-fsspec/flytekitplugins/fsspec/__init__.py | 5 ++++- plugins/flytekit-data-fsspec/setup.py | 1 + 4 files changed, 8 insertions(+), 2 deletions(-) diff --git a/flytekit/types/structured/basic_dfs.py b/flytekit/types/structured/basic_dfs.py index ff9d692cec..8113db5658 100644 --- a/flytekit/types/structured/basic_dfs.py +++ b/flytekit/types/structured/basic_dfs.py @@ -12,6 +12,7 @@ from flytekit.models.literals import StructuredDatasetMetadata from flytekit.models.types import StructuredDatasetType from flytekit.types.structured.structured_dataset import ( + GCS, LOCAL, PARQUET, S3, @@ -104,7 +105,7 @@ def decode( return pq.read_table(local_dir) -for protocol in [LOCAL, S3]: # Should we add GCS +for protocol in [LOCAL, S3, GCS]: StructuredDatasetTransformerEngine.register(PandasToParquetEncodingHandler(protocol), default_for_type=True) StructuredDatasetTransformerEngine.register(ParquetToPandasDecodingHandler(protocol), default_for_type=True) StructuredDatasetTransformerEngine.register(ArrowToParquetEncodingHandler(protocol), default_for_type=True) diff --git a/flytekit/types/structured/structured_dataset.py b/flytekit/types/structured/structured_dataset.py index fa57d7b5de..7d9bdcb9ba 100644 --- a/flytekit/types/structured/structured_dataset.py +++ b/flytekit/types/structured/structured_dataset.py @@ -37,6 +37,7 @@ # Protocols BIGQUERY = "bq" S3 = "s3" +GCS = "gs" LOCAL = "/" # For specifying the storage formats of StructuredDatasets. It's just a string, nothing fancy. diff --git a/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/__init__.py b/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/__init__.py index 85a65d17b8..4ec7c6e6ca 100644 --- a/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/__init__.py +++ b/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/__init__.py @@ -1,7 +1,7 @@ import importlib from flytekit import USE_STRUCTURED_DATASET, StructuredDatasetTransformerEngine, logger -from flytekit.types.structured.structured_dataset import S3 +from flytekit.types.structured.structured_dataset import GCS, S3 from .persist import FSSpecPersistence @@ -18,3 +18,6 @@ def _register(protocol: str): if importlib.util.find_spec("s3fs"): _register(S3) + + if importlib.util.find_spec("gcsfs"): + _register(GCS) diff --git a/plugins/flytekit-data-fsspec/setup.py b/plugins/flytekit-data-fsspec/setup.py index a0986a2e82..a20cc8a550 100644 --- a/plugins/flytekit-data-fsspec/setup.py +++ b/plugins/flytekit-data-fsspec/setup.py @@ -23,6 +23,7 @@ extras_require={ # https://github.com/fsspec/filesystem_spec/blob/master/setup.py#L36 "aws": ["s3fs>=2021.7.0"], + "gcp": ["gcsfs>=2021.7.0"], }, license="apache2", python_requires=">=3.7", From 2dd87381009eee4cdf6c188cec18336dc07d242f Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Wed, 2 Mar 2022 10:08:30 -0800 Subject: [PATCH 101/128] Make fetched entities callable within workflows (#867) Signed-off-by: Yee Hing Tong Signed-off-by: maximsmol --- flytekit/core/promise.py | 117 ++++++++++++++++-- flytekit/core/python_function_task.py | 7 +- flytekit/remote/launch_plan.py | 45 ++----- flytekit/remote/remote_callable.py | 61 +++++++++ flytekit/remote/task.py | 45 +------ flytekit/remote/workflow.py | 21 ++-- flytekit/tools/translator.py | 69 +++++++++++ .../flytekit/unit/core/test_serialization.py | 4 +- tests/flytekit/unit/remote/test_calling.py | 110 ++++++++++++++-- 9 files changed, 378 insertions(+), 101 deletions(-) create mode 100644 flytekit/remote/remote_callable.py diff --git a/flytekit/core/promise.py b/flytekit/core/promise.py index 17a4837432..a182ead2cc 100644 --- a/flytekit/core/promise.py +++ b/flytekit/core/promise.py @@ -548,11 +548,21 @@ def __rshift__(self, other: Any): return Output(*promises) # type: ignore +def binding_from_flyte_std( + ctx: _flyte_context.FlyteContext, + var_name: str, + expected_literal_type: _type_models.LiteralType, + t_value: typing.Any, +) -> _literals_models.Binding: + binding_data = binding_data_from_python_std(ctx, expected_literal_type, t_value, t_value_type=None) + return _literals_models.Binding(var=var_name, binding=binding_data) + + def binding_data_from_python_std( ctx: _flyte_context.FlyteContext, expected_literal_type: _type_models.LiteralType, t_value: typing.Any, - t_value_type: type, + t_value_type: Optional[type] = None, ) -> _literals_models.BindingData: # This handles the case where the given value is the output of another task if isinstance(t_value, Promise): @@ -568,7 +578,7 @@ def binding_data_from_python_std( if expected_literal_type.collection_type is None: raise AssertionError(f"this should be a list and it is not: {type(t_value)} vs {expected_literal_type}") - sub_type = ListTransformer.get_sub_type(t_value_type) + sub_type = ListTransformer.get_sub_type(t_value_type) if t_value_type else None collection = _literals_models.BindingDataCollection( bindings=[ binding_data_from_python_std(ctx, expected_literal_type.collection_type, t, sub_type) for t in t_value @@ -585,11 +595,11 @@ def binding_data_from_python_std( raise AssertionError( f"this should be a Dictionary type and it is not: {type(t_value)} vs {expected_literal_type}" ) - k_type, v_type = DictTransformer.get_dict_types(t_value_type) if expected_literal_type.simple == _type_models.SimpleType.STRUCT: lit = TypeEngine.to_literal(ctx, t_value, type(t_value), expected_literal_type) return _literals_models.BindingData(scalar=lit.scalar) else: + _, v_type = DictTransformer.get_dict_types(t_value_type) if t_value_type else None, None m = _literals_models.BindingDataMap( bindings={ k: binding_data_from_python_std(ctx, expected_literal_type.map_value_type, v, v_type) @@ -607,7 +617,7 @@ def binding_data_from_python_std( ) # This is the scalar case - e.g. my_task(in1=5) - scalar = TypeEngine.to_literal(ctx, t_value, t_value_type, expected_literal_type).scalar + scalar = TypeEngine.to_literal(ctx, t_value, t_value_type or type(t_value), expected_literal_type).scalar return _literals_models.BindingData(scalar=scalar) @@ -703,7 +713,8 @@ def __init__(self, node: Node, var: str): @property def node_id(self): """ - Override the underlying node_id property to refer to SdkNode. + Override the underlying node_id property to refer to the Node's id. This is to make sure that overriding + node IDs from with_overrides gets serialized correctly. :rtype: Text """ return self.node.id @@ -731,6 +742,19 @@ def construct_node_metadata(self) -> _workflow_model.NodeMetadata: ... +class HasFlyteInterface(Protocol): + @property + def name(self) -> str: + ... + + @property + def interface(self) -> _interface_models.TypedInterface: + ... + + def construct_node_metadata(self) -> _workflow_model.NodeMetadata: + ... + + def extract_obj_name(name: str) -> str: """ Generates a shortened name, without the module information. Useful for node-names etc. Only extracts the final @@ -743,6 +767,87 @@ def extract_obj_name(name: str) -> str: return name +def create_and_link_node_from_remote( + ctx: FlyteContext, + entity: HasFlyteInterface, + **kwargs, +): + """ + This method is used to generate a node with bindings. This is not used in the execution path. + """ + if ctx.compilation_state is None: + raise _user_exceptions.FlyteAssertion("Cannot create node when not compiling...") + + used_inputs = set() + bindings = [] + + typed_interface = entity.interface + + for k in sorted(typed_interface.inputs): + var = typed_interface.inputs[k] + if k not in kwargs: + raise _user_exceptions.FlyteAssertion("Input was not specified for: {} of type {}".format(k, var.type)) + v = kwargs[k] + # This check ensures that tuples are not passed into a function, as tuples are not supported by Flyte + # Usually a Tuple will indicate that multiple outputs from a previous task were accidentally passed + # into the function. + if isinstance(v, tuple): + raise AssertionError( + f"Variable({k}) for function({entity.name}) cannot receive a multi-valued tuple {v}." + f" Check if the predecessor function returning more than one value?" + ) + try: + bindings.append( + binding_from_flyte_std( + ctx, + var_name=k, + expected_literal_type=var.type, + t_value=v, + ) + ) + used_inputs.add(k) + except Exception as e: + raise AssertionError(f"Failed to Bind variable {k} for function {entity.name}.") from e + + extra_inputs = used_inputs ^ set(kwargs.keys()) + if len(extra_inputs) > 0: + raise _user_exceptions.FlyteAssertion( + "Too many inputs were specified for the interface. Extra inputs were: {}".format(extra_inputs) + ) + + # Detect upstream nodes + # These will be our core Nodes until we can amend the Promise to use NodeOutputs that reference our Nodes + upstream_nodes = list( + set( + [ + input_val.ref.node + for input_val in kwargs.values() + if isinstance(input_val, Promise) and input_val.ref.node_id != _common_constants.GLOBAL_INPUT_NODE_ID + ] + ) + ) + + flytekit_node = Node( + # TODO: Better naming, probably a derivative of the function name. + id=f"{ctx.compilation_state.prefix}n{len(ctx.compilation_state.nodes)}", + metadata=entity.construct_node_metadata(), + bindings=sorted(bindings, key=lambda b: b.var), + upstream_nodes=upstream_nodes, + flyte_entity=entity, + ) + ctx.compilation_state.add_node(flytekit_node) + + if len(typed_interface.outputs) == 0: + return VoidPromise(entity.name) + + # Create a node output object for each output, they should all point to this node of course. + node_outputs = [] + for output_name, output_var_model in typed_interface.outputs.items(): + node_outputs.append(Promise(output_name, NodeOutput(node=flytekit_node, var=output_name))) + + return create_task_output(node_outputs) + + def create_and_link_node( ctx: FlyteContext, entity: SupportsNodeCreation, @@ -819,8 +924,6 @@ def create_and_link_node( # Create a node output object for each output, they should all point to this node of course. node_outputs = [] for output_name, output_var_model in typed_interface.outputs.items(): - # TODO: If node id gets updated later, we have to make sure to update the NodeOutput model's ID, which - # is currently just a static str node_outputs.append(Promise(output_name, NodeOutput(node=flytekit_node, var=output_name))) # Don't print this, it'll crash cuz sdk_node._upstream_node_ids might be None, but idl code will break diff --git a/flytekit/core/python_function_task.py b/flytekit/core/python_function_task.py index 35016f9571..2e102c8b62 100644 --- a/flytekit/core/python_function_task.py +++ b/flytekit/core/python_function_task.py @@ -210,7 +210,12 @@ def compile_into_workflow( for entity, model in model_entities.items(): # We only care about gathering tasks here. Launch plans are handled by # propeller. Subworkflows should already be in the workflow spec. - if not isinstance(entity, Task): + if not isinstance(entity, Task) and not isinstance(entity, task_models.TaskTemplate): + continue + + # Handle FlyteTask + if isinstance(entity, task_models.TaskTemplate): + tts.append(entity) continue # We are currently not supporting reference tasks since these will diff --git a/flytekit/remote/launch_plan.py b/flytekit/remote/launch_plan.py index 3bf845fcab..32b3a93aec 100644 --- a/flytekit/remote/launch_plan.py +++ b/flytekit/remote/launch_plan.py @@ -1,60 +1,32 @@ +from __future__ import annotations + from typing import Optional +from flytekit.core import hash as hash_mixin from flytekit.core.interface import Interface -from flytekit.core.launch_plan import ReferenceLaunchPlan -from flytekit.core.type_engine import TypeEngine -from flytekit.loggers import remote_logger as logger from flytekit.models import interface as _interface_models from flytekit.models import launch_plan as _launch_plan_models from flytekit.models.core import identifier as id_models from flytekit.remote import interface as _interface +from flytekit.remote.remote_callable import RemoteEntity -class FlyteLaunchPlan(_launch_plan_models.LaunchPlanSpec): +class FlyteLaunchPlan(hash_mixin.HashOnReferenceMixin, RemoteEntity, _launch_plan_models.LaunchPlanSpec): """A class encapsulating a remote Flyte launch plan.""" def __init__(self, id, *args, **kwargs): super(FlyteLaunchPlan, self).__init__(*args, **kwargs) # Set all the attributes we expect this class to have self._id = id + self._name = id.name # The interface is not set explicitly unless fetched in an engine context self._interface = None self._python_interface = None - self._reference_entity = None - - def __call__(self, *args, **kwargs): - if self.reference_entity is None: - logger.warning( - f"FlyteLaunchPlan {self} is not callable, most likely because flytekit could not " - f"guess the python interface. The workflow calling this launch plan may not behave correctly." - ) - return - return self.reference_entity(*args, **kwargs) - # TODO: Refactor behind mixin @property - def reference_entity(self) -> Optional[ReferenceLaunchPlan]: - if self._reference_entity is None: - if self.guessed_python_interface is None: - try: - self.guessed_python_interface = Interface( - TypeEngine.guess_python_types(self.interface.inputs), - TypeEngine.guess_python_types(self.interface.outputs), - ) - except Exception as e: - logger.warning(f"Error backing out interface {e}, Flyte interface {self.interface}") - return None - - self._reference_entity = ReferenceLaunchPlan( - self.id.project, - self.id.domain, - self.id.name, - self.id.version, - inputs=self.guessed_python_interface.inputs, - outputs=self.guessed_python_interface.outputs, - ) - return self._reference_entity + def name(self) -> str: + return self._name @classmethod def promote_from_model( @@ -71,7 +43,6 @@ def promote_from_model( auth_role=model.auth_role, raw_output_data_config=model.raw_output_data_config, ) - return lp @property diff --git a/flytekit/remote/remote_callable.py b/flytekit/remote/remote_callable.py new file mode 100644 index 0000000000..699d4d8140 --- /dev/null +++ b/flytekit/remote/remote_callable.py @@ -0,0 +1,61 @@ +from abc import ABC, abstractmethod +from typing import Any, Optional, Tuple, Union + +from flytekit.core.context_manager import BranchEvalMode, ExecutionState, FlyteContext +from flytekit.core.promise import Promise, VoidPromise, create_and_link_node_from_remote, extract_obj_name +from flytekit.exceptions import user as user_exceptions +from flytekit.loggers import remote_logger as logger +from flytekit.models.core.workflow import NodeMetadata + + +class RemoteEntity(ABC): + @property + @abstractmethod + def name(self) -> str: + ... + + def construct_node_metadata(self) -> NodeMetadata: + """ + Used when constructing the node that encapsulates this task as part of a broader workflow definition. + """ + return NodeMetadata( + name=extract_obj_name(self.name), + ) + + def compile(self, ctx: FlyteContext, *args, **kwargs): + return create_and_link_node_from_remote(ctx, entity=self, **kwargs) # noqa + + def __call__(self, *args, **kwargs): + # When a Task is () aka __called__, there are three things we may do: + # a. Plain execution Mode - just run the execute function. If not overridden, we should raise an exception + # b. Compilation Mode - this happens when the function is called as part of a workflow (potentially + # dynamic task). Produce promise objects and create a node. + # c. Workflow Execution Mode - when a workflow is being run locally. Even though workflows are functions + # and everything should be able to be passed through naturally, we'll want to wrap output values of the + # function into objects, so that potential .with_cpu or other ancillary functions can be attached to do + # nothing. Subsequent tasks will have to know how to unwrap these. If by chance a non-Flyte task uses a + # task output as an input, things probably will fail pretty obviously. + # Since this is a reference entity, it still needs to be mocked otherwise an exception will be raised. + if len(args) > 0: + raise user_exceptions.FlyteAssertion( + f"Cannot call remotely fetched entity with args - detected {len(args)} positional args {args}" + ) + + ctx = FlyteContext.current_context() + if ctx.compilation_state is not None and ctx.compilation_state.mode == 1: + return self.compile(ctx, *args, **kwargs) + elif ( + ctx.execution_state is not None and ctx.execution_state.mode == ExecutionState.Mode.LOCAL_WORKFLOW_EXECUTION + ): + if ctx.execution_state.branch_eval_mode == BranchEvalMode.BRANCH_SKIPPED: + return + return self.local_execute(ctx, **kwargs) + else: + logger.debug("Fetched entity, running raw execute.") + return self.execute(**kwargs) + + def local_execute(self, ctx: FlyteContext, **kwargs) -> Optional[Union[Tuple[Promise], Promise, VoidPromise]]: + raise Exception("Remotely fetched entities cannot be run locally. You have to mock this out.") + + def execute(self, **kwargs) -> Any: + raise Exception("Remotely fetched entities cannot be run locally. You have to mock this out.") diff --git a/flytekit/remote/task.py b/flytekit/remote/task.py index 34b4f1d9c2..028719ec7f 100644 --- a/flytekit/remote/task.py +++ b/flytekit/remote/task.py @@ -1,16 +1,16 @@ from typing import Optional -from flytekit.core import hash as _hash_mixin +from flytekit.core import hash as hash_mixin from flytekit.core.interface import Interface -from flytekit.core.task import ReferenceTask from flytekit.core.type_engine import TypeEngine from flytekit.loggers import remote_logger as logger from flytekit.models import task as _task_model from flytekit.models.core import identifier as _identifier_model from flytekit.remote import interface as _interfaces +from flytekit.remote.remote_callable import RemoteEntity -class FlyteTask(_hash_mixin.HashOnReferenceMixin, _task_model.TaskTemplate): +class FlyteTask(hash_mixin.HashOnReferenceMixin, RemoteEntity, _task_model.TaskTemplate): """A class encapsulating a remote Flyte task.""" def __init__(self, id, type, metadata, interface, custom, container=None, task_type_version=0, config=None): @@ -25,44 +25,11 @@ def __init__(self, id, type, metadata, interface, custom, container=None, task_t config=config, ) self._python_interface = None - self._reference_entity = None - - def __call__(self, *args, **kwargs): - if self.reference_entity is None: - logger.warning( - f"FlyteTask {self} is not callable, most likely because flytekit could not " - f"guess the python interface. The workflow calling this task may not behave correctly" - ) - return - return self.reference_entity(*args, **kwargs) - - # TODO: Refactor behind mixin - @property - def reference_entity(self) -> Optional[ReferenceTask]: - if self._reference_entity is None: - if self.guessed_python_interface is None: - try: - self.guessed_python_interface = Interface( - TypeEngine.guess_python_types(self.interface.inputs), - TypeEngine.guess_python_types(self.interface.outputs), - ) - except Exception as e: - logger.warning(f"Error backing out interface {e}, Flyte interface {self.interface}") - return None - - self._reference_entity = ReferenceTask( - self.id.project, - self.id.domain, - self.id.name, - self.id.version, - inputs=self.guessed_python_interface.inputs, - outputs=self.guessed_python_interface.outputs, - ) - return self._reference_entity + self._name = id.name @property - def interface(self) -> _interfaces.TypedInterface: - return super(FlyteTask, self).interface + def name(self) -> str: + return self._name @property def resource_type(self) -> _identifier_model.ResourceType: diff --git a/flytekit/remote/workflow.py b/flytekit/remote/workflow.py index 14ef6c91bf..2ff74376c8 100644 --- a/flytekit/remote/workflow.py +++ b/flytekit/remote/workflow.py @@ -14,21 +14,22 @@ from flytekit.models.core import workflow as _workflow_models from flytekit.remote import interface as _interfaces from flytekit.remote import nodes as _nodes +from flytekit.remote.remote_callable import RemoteEntity -class FlyteWorkflow(_hash_mixin.HashOnReferenceMixin, _workflow_models.WorkflowTemplate): +class FlyteWorkflow(_hash_mixin.HashOnReferenceMixin, RemoteEntity, _workflow_models.WorkflowTemplate): """A class encapsulating a remote Flyte workflow.""" def __init__( self, + id: id_models.Identifier, nodes: List[_nodes.FlyteNode], interface, output_bindings, - id: id_models.Identifier, metadata, metadata_defaults, subworkflows: Optional[Dict[id_models.Identifier, _workflow_models.WorkflowTemplate]] = None, - tasks: Optional[Dict[id_models.Identifier, _task_models.TaskSpec]] = None, + tasks: Optional[Dict[id_models.Identifier, _task_models.TaskTemplate]] = None, launch_plans: Optional[Dict[id_models.Identifier, launch_plan_models.LaunchPlanSpec]] = None, compiled_closure: Optional[compiler_models.CompiledWorkflowClosure] = None, ): @@ -58,10 +59,15 @@ def __init__( self._launch_plans = launch_plans self._compiled_closure = compiled_closure self._node_map = None + self._name = id.name + + @property + def name(self) -> str: + return self._name @property - def interface(self) -> _interfaces.TypedInterface: - return super(FlyteWorkflow, self).interface + def sub_workflows(self) -> Optional[Dict[id_models.Identifier, _workflow_models.WorkflowTemplate]]: + return self._subworkflows @property def entity_type_text(self) -> str: @@ -114,8 +120,8 @@ def promote_from_model( # No inputs/outputs specified, see the constructor for more information on the overrides. wf = cls( - nodes=list(node_map.values()), id=base_model.id, + nodes=list(node_map.values()), metadata=base_model.metadata, metadata_defaults=base_model.metadata_defaults, interface=_interfaces.TypedInterface.promote_from_model(base_model.interface), @@ -159,6 +165,3 @@ def promote_from_closure( ) flyte_wf._compiled_closure = closure return flyte_wf - - def __call__(self, *args, **input_map): - raise NotImplementedError diff --git a/flytekit/tools/translator.py b/flytekit/tools/translator.py index 8c9750e8cd..05b4224cc9 100644 --- a/flytekit/tools/translator.py +++ b/flytekit/tools/translator.py @@ -126,6 +126,9 @@ def get_serializable_workflow( settings: SerializationSettings, entity: WorkflowBase, ) -> admin_workflow_models.WorkflowSpec: + # TODO: Try to move up following config refactor - https://github.com/flyteorg/flyte/issues/2214 + from flytekit.remote.workflow import FlyteWorkflow + # Get node models upstream_node_models = [ get_serializable(entity_mapping, settings, n) @@ -151,6 +154,11 @@ def get_serializable_workflow( sub_wfs.append(sub_wf_spec.template) sub_wfs.extend(sub_wf_spec.sub_workflows) + if isinstance(n.flyte_entity, FlyteWorkflow): + get_serializable(entity_mapping, settings, n.flyte_entity) + sub_wfs.append(n.flyte_entity) + sub_wfs.extend([s for s in n.flyte_entity.sub_workflows.values()]) + if isinstance(n.flyte_entity, BranchNode): if_else: workflow_model.IfElseBlock = n.flyte_entity._ifelse_block # See comment in get_serializable_branch_node also. Again this is a List[Node] even though it's supposed @@ -168,6 +176,10 @@ def get_serializable_workflow( sub_wf_spec = get_serializable(entity_mapping, settings, leaf_node.flyte_entity) sub_wfs.append(sub_wf_spec.template) sub_wfs.extend(sub_wf_spec.sub_workflows) + elif isinstance(leaf_node.flyte_entity, FlyteWorkflow): + get_serializable(entity_mapping, settings, leaf_node.flyte_entity) + sub_wfs.append(leaf_node.flyte_entity) + sub_wfs.extend([s for s in leaf_node.flyte_entity.sub_workflows.values()]) wf_id = _identifier_model.Identifier( resource_type=_identifier_model.ResourceType.WORKFLOW, @@ -237,6 +249,11 @@ def get_serializable_node( if entity.flyte_entity is None: raise Exception(f"Node {entity.id} has no flyte entity") + # TODO: Try to move back up following config refactor - https://github.com/flyteorg/flyte/issues/2214 + from flytekit.remote.launch_plan import FlyteLaunchPlan + from flytekit.remote.task import FlyteTask + from flytekit.remote.workflow import FlyteWorkflow + upstream_sdk_nodes = [ get_serializable(entity_mapping, settings, n) for n in entity.upstream_nodes @@ -319,6 +336,49 @@ def get_serializable_node( output_aliases=[], workflow_node=workflow_model.WorkflowNode(launchplan_ref=lp_spec.id), ) + + elif isinstance(entity.flyte_entity, FlyteTask): + # Recursive call doesn't do anything except put the entity on the map. + get_serializable(entity_mapping, settings, entity.flyte_entity) + node_model = workflow_model.Node( + id=_dnsify(entity.id), + metadata=entity.metadata, + inputs=entity.bindings, + upstream_node_ids=[n.id for n in upstream_sdk_nodes], + output_aliases=[], + task_node=workflow_model.TaskNode( + reference_id=entity.flyte_entity.id, overrides=TaskNodeOverrides(resources=entity._resources) + ), + ) + elif isinstance(entity.flyte_entity, FlyteWorkflow): + wf_template = get_serializable(entity_mapping, settings, entity.flyte_entity) + for _, sub_wf in entity.flyte_entity.sub_workflows.items(): + get_serializable(entity_mapping, settings, sub_wf) + node_model = workflow_model.Node( + id=_dnsify(entity.id), + metadata=entity.metadata, + inputs=entity.bindings, + upstream_node_ids=[n.id for n in upstream_sdk_nodes], + output_aliases=[], + workflow_node=workflow_model.WorkflowNode(sub_workflow_ref=wf_template.id), + ) + elif isinstance(entity.flyte_entity, FlyteLaunchPlan): + # Recursive call doesn't do anything except put the entity on the map. + get_serializable(entity_mapping, settings, entity.flyte_entity) + # Node's inputs should not contain the data which is fixed input + node_input = [] + for b in entity.bindings: + if b.var not in entity.flyte_entity.fixed_inputs.literals: + node_input.append(b) + + node_model = workflow_model.Node( + id=_dnsify(entity.id), + metadata=entity.metadata, + inputs=node_input, + upstream_node_ids=[n.id for n in upstream_sdk_nodes], + output_aliases=[], + workflow_node=workflow_model.WorkflowNode(launchplan_ref=entity.flyte_entity.id), + ) else: raise Exception(f"Node contained non-serializable entity {entity._flyte_entity}") @@ -375,6 +435,11 @@ def get_serializable( :return: The resulting control plane entity, in addition to being added to the mutable entity_mapping parameter is also returned. """ + # TODO: Try to replace following config refactor - https://github.com/flyteorg/flyte/issues/2214 + from flytekit.remote.launch_plan import FlyteLaunchPlan + from flytekit.remote.task import FlyteTask + from flytekit.remote.workflow import FlyteWorkflow + if entity in entity_mapping: return entity_mapping[entity] @@ -395,6 +460,10 @@ def get_serializable( elif isinstance(entity, BranchNode): cp_entity = get_serializable_branch_node(entity_mapping, settings, entity) + + elif isinstance(entity, FlyteTask) or isinstance(entity, FlyteWorkflow) or isinstance(entity, FlyteLaunchPlan): + cp_entity = entity + else: raise Exception(f"Non serializable type found {type(entity)} Entity {entity}") diff --git a/tests/flytekit/unit/core/test_serialization.py b/tests/flytekit/unit/core/test_serialization.py index d3395e9fd5..11616203cd 100644 --- a/tests/flytekit/unit/core/test_serialization.py +++ b/tests/flytekit/unit/core/test_serialization.py @@ -354,7 +354,7 @@ def middle_subwf() -> typing.Tuple[int, int]: @workflow def parent_wf() -> typing.Tuple[int, int, int, int]: m1, m2 = middle_subwf() - l1, l2 = leaf_subwf() + l1, l2 = leaf_subwf().with_overrides(node_name="foo-node") return m1, m2, l1, l2 wf_spec = get_serializable(OrderedDict(), serialization_settings, parent_wf) @@ -366,6 +366,8 @@ def parent_wf() -> typing.Tuple[int, int, int, int]: assert len(midwf.nodes) == 1 assert midwf.nodes[0].workflow_node is not None assert midwf.nodes[0].workflow_node.sub_workflow_ref.name == "test_serialization.leaf_subwf" + assert wf_spec.template.nodes[1].id == "foo-node" + assert wf_spec.template.outputs[2].binding.promise.node_id == "foo-node" def test_serialization_named_outputs_single(): diff --git a/tests/flytekit/unit/remote/test_calling.py b/tests/flytekit/unit/remote/test_calling.py index 97a1a001dc..07483487ab 100644 --- a/tests/flytekit/unit/remote/test_calling.py +++ b/tests/flytekit/unit/remote/test_calling.py @@ -3,14 +3,19 @@ import pytest +from flytekit import dynamic from flytekit.core import context_manager -from flytekit.core.context_manager import Image, ImageConfig +from flytekit.core.context_manager import ExecutionState, FastSerializationSettings, Image, ImageConfig from flytekit.core.launch_plan import LaunchPlan -from flytekit.core.reference_entity import ReferenceSpec from flytekit.core.task import task +from flytekit.core.type_engine import TypeEngine from flytekit.core.workflow import workflow +from flytekit.exceptions.user import FlyteAssertion +from flytekit.models.core.workflow import WorkflowTemplate +from flytekit.models.task import TaskTemplate from flytekit.remote import FlyteLaunchPlan, FlyteTask from flytekit.remote.interface import TypedInterface +from flytekit.remote.workflow import FlyteWorkflow from flytekit.tools.translator import gather_dependent_entities, get_serializable default_img = Image(name="default", fqn="test", tag="tag") @@ -48,18 +53,28 @@ def sub_wf(a: int, b: str) -> (int, str): def test_fetched_task(): @workflow def wf(a: int) -> int: - return ft(a=a) + return ft(a=a).with_overrides(node_name="foobar") # Should not work unless mocked out. with pytest.raises(Exception, match="cannot be run locally"): wf(a=3) - # Should have one reference entity + # Should have one task template serialized = OrderedDict() - get_serializable(serialized, serialization_settings, wf) + wf_spec = get_serializable(serialized, serialization_settings, wf) vals = [v for v in serialized.values()] - refs = [f for f in filter(lambda x: isinstance(x, ReferenceSpec), vals)] - assert len(refs) == 1 + tts = [f for f in filter(lambda x: isinstance(x, TaskTemplate), vals)] + assert len(tts) == 1 + assert wf_spec.template.nodes[0].id == "foobar" + assert wf_spec.template.outputs[0].binding.promise.node_id == "foobar" + + +def test_misnamed(): + with pytest.raises(FlyteAssertion): + + @workflow + def wf(a: int) -> int: + return ft(b=a) def test_calling_lp(): @@ -83,3 +98,84 @@ def wf2(a: int) -> typing.Tuple[int, str]: wf_spec = get_serializable(serialized, serialization_settings, wf2) print(wf_spec.template.nodes[0].workflow_node.launchplan_ref) assert wf_spec.template.nodes[0].workflow_node.launchplan_ref == lp_model.id + + +def test_dynamic(): + @dynamic + def my_subwf(a: int) -> typing.List[int]: + s = [] + for i in range(a): + s.append(ft(a=i)) + return s + + with context_manager.FlyteContextManager.with_context( + context_manager.FlyteContextManager.current_context().with_serialization_settings( + context_manager.SerializationSettings( + project="test_proj", + domain="test_domain", + version="abc", + image_config=ImageConfig(Image(name="name", fqn="image", tag="name")), + env={}, + fast_serialization_settings=FastSerializationSettings(enabled=True), + ) + ) + ) as ctx: + with context_manager.FlyteContextManager.with_context( + ctx.with_execution_state( + ctx.execution_state.with_params( + mode=ExecutionState.Mode.TASK_EXECUTION, + additional_context={ + "dynamic_addl_distro": "s3://my-s3-bucket/fast/123", + "dynamic_dest_dir": "/User/flyte/workflows", + }, + ) + ) + ) as ctx: + input_literal_map = TypeEngine.dict_to_literal_map(ctx, {"a": 2}) + # Test that it works + dynamic_job_spec = my_subwf.dispatch_execute(ctx, input_literal_map) + assert len(dynamic_job_spec._nodes) == 2 + assert len(dynamic_job_spec.tasks) == 1 + assert dynamic_job_spec.tasks[0].id == ft.id + + # Test that the fast execute stuff does not get applied because the commands of tasks fetched from + # Admin should never change. + args = " ".join(dynamic_job_spec.tasks[0].container.args) + assert not args.startswith("pyflyte-fast-execute") + + +def test_calling_wf(): + # No way to fetch from Admin in unit tests so we serialize and then promote back + serialized = OrderedDict() + wf_spec = get_serializable(serialized, serialization_settings, sub_wf) + task_templates, wf_specs, lp_specs = gather_dependent_entities(serialized) + fwf = FlyteWorkflow.promote_from_model(wf_spec.template, tasks=task_templates) + + @workflow + def parent_1(a: int, b: str) -> typing.Tuple[int, str]: + y = t1(a=a) + return fwf(a=y, b=b) + + # No way to fetch from Admin in unit tests so we serialize and then promote back + serialized = OrderedDict() + wf_spec = get_serializable(serialized, serialization_settings, parent_1) + # Get task_specs from the second one, merge with the first one. Admin normally would be the one to do this. + task_templates_p1, wf_specs, lp_specs = gather_dependent_entities(serialized) + for k, v in task_templates.items(): + task_templates_p1[k] = v + + # Pick out the subworkflow templates from the ordereddict. We can't use the output of the gather_dependent_entities + # function because that only looks for WorkflowSpecs + subwf_templates = {x.id: x for x in list(filter(lambda x: isinstance(x, WorkflowTemplate), serialized.values()))} + fwf_p1 = FlyteWorkflow.promote_from_model(wf_spec.template, sub_workflows=subwf_templates, tasks=task_templates_p1) + + @workflow + def parent_2(a: int, b: str) -> typing.Tuple[int, str]: + x, y = fwf_p1(a=a, b=b) + z = t1(a=x) + return z, y + + serialized = OrderedDict() + wf_spec = get_serializable(serialized, serialization_settings, parent_2) + # Make sure both were picked up. + assert len(wf_spec.sub_workflows) == 2 From 93b0fd23e9677f4d8777e00b532b59b81b1c5903 Mon Sep 17 00:00:00 2001 From: Eduardo Apolinario <653394+eapolinario@users.noreply.github.com> Date: Wed, 2 Mar 2022 13:24:45 -0800 Subject: [PATCH 102/128] Caching of offloaded objects (#762) * Remove flyteidl from install_requires Signed-off-by: Eduardo Apolinario * Expose hash in Literal Signed-off-by: Eduardo Apolinario * Set hash in TypeEngine Signed-off-by: Eduardo Apolinario * Modify cache key calculation to take hash into account Signed-off-by: Eduardo Apolinario * Opt-in PandasDataFrameTransformer Signed-off-by: Eduardo Apolinario * Add unit tests Signed-off-by: Eduardo Apolinario * Iterate using a flyteidl branch Signed-off-by: Eduardo Apolinario * Regenerate requirements files Signed-off-by: Eduardo Apolinario * Regenerate requirements files Signed-off-by: Eduardo Apolinario * Move _hash_overridable to StructureDatasetTransformerEngine Signed-off-by: Eduardo Apolinario * Move HashMethod to flytekit.core.hash Signed-off-by: Eduardo Apolinario * Fix `unit_test` make target Signed-off-by: Eduardo Apolinario * Split `unit_test` make target in two lines Signed-off-by: Eduardo Apolinario * Add assert to structured dataset compatibility test Signed-off-by: Eduardo Apolinario * Remove TODO Signed-off-by: Eduardo Apolinario * Regenerate plugins requirements files pointing to the right version of flyteidl. Signed-off-by: Eduardo Apolinario * Set hash as a property of the literal Signed-off-by: Eduardo Apolinario * Install plugins requirements in CI. Signed-off-by: Eduardo Apolinario * Add hash.setter Signed-off-by: Eduardo Apolinario * Install flyteidl directly Signed-off-by: Eduardo Apolinario * Revert "Regenerate plugins requirements files pointing to the right version of flyteidl." This reverts commit c2dbb540ad716e883932dd9255045c128e16c6a0. Signed-off-by: Eduardo Apolinario * wip - Add support for univariate lists Signed-off-by: Eduardo Apolinario * Add support for lists of annotated objects Signed-off-by: Eduardo Apolinario * Revamp generation of cache key (to cover case of literals collections and maps) Signed-off-by: Eduardo Apolinario * Leave TODO for warning Signed-off-by: Eduardo Apolinario * Revert "Add support for lists of annotated objects" This reverts commit 4b5f608f762a018eae1f42acc3bcf6c4cf12583a. Signed-off-by: Eduardo Apolinario * Revert "wip - Add support for univariate lists" This reverts commit adaa44884f76bd4c7c2225e8e309132cb9111906. Signed-off-by: Eduardo Apolinario * Remove docstring Signed-off-by: Eduardo Apolinario * Add flyteidl>=0.23.0 Signed-off-by: Eduardo Apolinario * Remove mentions to branch flyteidl@add-hash-to-literal Signed-off-by: Eduardo Apolinario * Bump flyteidl in plugins requirements Signed-off-by: Eduardo Apolinario * Regenerate plugins requirements again Signed-off-by: Eduardo Apolinario * Restore papermill/requirements.txt Signed-off-by: Eduardo Apolinario * Point flytekitplugins-spark to the offloaded-objects-caching branch in papermill tests Signed-off-by: Eduardo Apolinario * Set flyteidl>=0.23.0 in papermill dev-requirements Co-authored-by: Eduardo Apolinario --- Makefile | 3 +- dev-requirements.txt | 22 ++- doc-requirements.txt | 19 ++- flytekit/__init__.py | 1 + flytekit/core/hash.py | 20 +++ flytekit/core/local_cache.py | 30 +++- flytekit/core/type_engine.py | 52 ++++--- flytekit/models/literals.py | 19 ++- flytekit/types/schema/types_pandas.py | 3 + .../types/structured/structured_dataset.py | 3 + plugins/flytekit-aws-athena/requirements.txt | 40 +++-- plugins/flytekit-aws-batch/requirements.txt | 62 +++++--- .../flytekit-aws-sagemaker/requirements.txt | 50 ++++--- plugins/flytekit-bigquery/requirements.txt | 40 ++--- plugins/flytekit-data-fsspec/requirements.txt | 48 ++++-- plugins/flytekit-dolt/requirements.txt | 40 +++-- .../requirements.txt | 104 +++++++------ plugins/flytekit-hive/requirements.txt | 40 +++-- plugins/flytekit-k8s-pod/requirements.txt | 44 ++++-- plugins/flytekit-kf-mpi/requirements.txt | 40 +++-- plugins/flytekit-kf-pytorch/requirements.txt | 40 +++-- .../flytekit-kf-tensorflow/requirements.txt | 40 +++-- plugins/flytekit-pandera/requirements.txt | 45 ++++-- .../flytekit-papermill/dev-requirements.in | 3 +- .../flytekit-papermill/dev-requirements.txt | 26 +++- plugins/flytekit-snowflake/requirements.txt | 40 +++-- plugins/flytekit-spark/requirements.txt | 40 +++-- plugins/flytekit-sqlalchemy/requirements.txt | 40 +++-- requirements-spark2.txt | 18 ++- requirements.txt | 22 ++- .../workflows/requirements.txt | 4 +- tests/flytekit/unit/core/test_local_cache.py | 139 ++++++++++++++++-- tests/flytekit/unit/core/test_type_engine.py | 76 ++++++++++ tests/flytekit/unit/core/test_type_hints.py | 109 ++++++++++++++ .../test_structured_dataset.py | 4 + 35 files changed, 1000 insertions(+), 326 deletions(-) diff --git a/Makefile b/Makefile index b5464e205e..404c1eb1fd 100644 --- a/Makefile +++ b/Makefile @@ -50,7 +50,8 @@ test: lint unit_test .PHONY: unit_test unit_test: - FLYTE_SDK_USE_STRUCTURED_DATASET=TRUE pytest tests/flytekit/unit tests/flytekit_compatibility + FLYTE_SDK_USE_STRUCTURED_DATASET=FALSE pytest tests/flytekit_compatibility && \ + FLYTE_SDK_USE_STRUCTURED_DATASET=TRUE pytest tests/flytekit/unit requirements-spark2.txt: export CUSTOM_COMPILE_COMMAND := make requirements-spark2.txt requirements-spark2.txt: requirements-spark2.in install-piptools diff --git a/dev-requirements.txt b/dev-requirements.txt index 65e61b91fd..af1217527e 100644 --- a/dev-requirements.txt +++ b/dev-requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.8 +# This file is autogenerated by pip-compile with python 3.9 # To update, run: # # make dev-requirements.txt @@ -72,7 +72,10 @@ croniter==1.3.4 # -c requirements.txt # flytekit cryptography==36.0.1 - # via paramiko + # via + # -c requirements.txt + # paramiko + # secretstorage dataclasses-json==0.5.6 # via # -c requirements.txt @@ -113,7 +116,7 @@ docstring-parser==0.13 # flytekit filelock==3.6.0 # via virtualenv -flyteidl==0.22.3 +flyteidl==0.23.0 # via # -c requirements.txt # flytekit @@ -163,6 +166,11 @@ importlib-metadata==4.11.2 # keyring iniconfig==1.1.1 # via pytest +jeepney==0.7.1 + # via + # -c requirements.txt + # keyring + # secretstorage jinja2==3.0.3 # via # -c requirements.txt @@ -275,7 +283,9 @@ pyasn1==0.4.8 pyasn1-modules==0.2.8 # via google-auth pycparser==2.21 - # via cffi + # via + # -c requirements.txt + # cffi pynacl==1.5.0 # via paramiko pyparsing==3.0.7 @@ -349,6 +359,10 @@ retry==0.9.2 # flytekit rsa==4.8 # via google-auth +secretstorage==3.3.1 + # via + # -c requirements.txt + # keyring six==1.16.0 # via # -c requirements.txt diff --git a/doc-requirements.txt b/doc-requirements.txt index f734be7406..c4c496447d 100644 --- a/doc-requirements.txt +++ b/doc-requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.8 +# This file is autogenerated by pip-compile with python 3.9 # To update, run: # # make doc-requirements.txt @@ -42,7 +42,9 @@ cookiecutter==1.7.3 croniter==1.3.4 # via flytekit cryptography==36.0.1 - # via -r doc-requirements.in + # via + # -r doc-requirements.in + # secretstorage css-html-js-minify==2.5.5 # via sphinx-material dataclasses-json==0.5.6 @@ -61,7 +63,7 @@ docutils==0.17.1 # via # sphinx # sphinx-panels -flyteidl==0.22.3 +flyteidl==0.23.0 # via flytekit furo @ git+https://github.com/flyteorg/furo@main # via -r doc-requirements.in @@ -79,6 +81,10 @@ importlib-metadata==4.11.2 # via # keyring # sphinx +jeepney==0.7.1 + # via + # keyring + # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -125,10 +131,7 @@ protobuf==3.19.4 # googleapis-common-protos # protoc-gen-swagger protoc-gen-swagger==0.1.0 - # via - # flyteidl - # flytekit - + # via flyteidl py==1.11.0 # via retry pyarrow==6.0.1 @@ -174,6 +177,8 @@ responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit +secretstorage==3.3.1 + # via keyring six==1.16.0 # via # cookiecutter diff --git a/flytekit/__init__.py b/flytekit/__init__.py index 6e2e614702..8089e8bd74 100644 --- a/flytekit/__init__.py +++ b/flytekit/__init__.py @@ -169,6 +169,7 @@ from flytekit.core.context_manager import ExecutionParameters, FlyteContext, FlyteContextManager from flytekit.core.data_persistence import DataPersistence, DataPersistencePlugins from flytekit.core.dynamic_workflow_task import dynamic +from flytekit.core.hash import HashMethod from flytekit.core.launch_plan import LaunchPlan from flytekit.core.map_task import map_task from flytekit.core.notification import Email, PagerDuty, Slack diff --git a/flytekit/core/hash.py b/flytekit/core/hash.py index 8b2559f696..2aca94e846 100644 --- a/flytekit/core/hash.py +++ b/flytekit/core/hash.py @@ -1,3 +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) diff --git a/flytekit/core/local_cache.py b/flytekit/core/local_cache.py index 9180224207..4afab73955 100644 --- a/flytekit/core/local_cache.py +++ b/flytekit/core/local_cache.py @@ -1,16 +1,42 @@ +import base64 from typing import Optional +import cloudpickle from diskcache import Cache -from flytekit.models.literals import LiteralMap +from flytekit.models.literals import Literal, LiteralCollection, LiteralMap # Location on the filesystem where serialized objects will be stored # TODO: read from config CACHE_LOCATION = "~/.flyte/local-cache" +def _recursive_hash_placement(literal: Literal) -> Literal: + if literal.collection is not None: + literals = [_recursive_hash_placement(literal) for literal in literal.collection.literals] + return Literal(collection=LiteralCollection(literals=literals)) + elif literal.map is not None: + literal_map = {} + for key, literal in literal.map.literals.items(): + literal_map[key] = _recursive_hash_placement(literal) + return Literal(map=LiteralMap(literal_map)) + + # Base case + if literal.hash is not None: + return Literal(hash=literal.hash) + else: + return literal + + def _calculate_cache_key(task_name: str, cache_version: str, input_literal_map: LiteralMap) -> str: - return f"{task_name}-{cache_version}-{input_literal_map}" + # Traverse the literals and replace the literal with a new literal that only contains the hash + literal_map_overridden = {} + for key, literal in input_literal_map.literals.items(): + literal_map_overridden[key] = _recursive_hash_placement(literal) + + # Pickle the literal map and use base64 encoding to generate a representation of it + b64_encoded = base64.b64encode(cloudpickle.dumps(LiteralMap(literal_map_overridden))) + return f"{task_name}-{cache_version}-{b64_encoded}" class LocalTaskCache(object): diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 7a405df6e3..c7ae0b3dd3 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -10,13 +10,6 @@ from abc import ABC, abstractmethod from typing import NamedTuple, Optional, Type, cast -from typing_extensions import get_args as _get_args - -try: - from typing import Annotated, get_args, get_origin -except ImportError: - from typing_extensions import Annotated, get_origin, get_args - from dataclasses_json import DataClassJsonMixin, dataclass_json from google.protobuf import json_format as _json_format from google.protobuf import reflection as _proto_reflection @@ -26,9 +19,11 @@ from google.protobuf.struct_pb2 import Struct from marshmallow_enum import EnumField, LoadDumpOptions from marshmallow_jsonschema import JSONSchema +from typing_extensions import Annotated, get_args, get_origin from flytekit.core.annotation import FlyteAnnotation from flytekit.core.context_manager import FlyteContext +from flytekit.core.hash import HashMethod from flytekit.core.type_helpers import load_type_from_tag from flytekit.exceptions import user as user_exceptions from flytekit.loggers import logger @@ -52,14 +47,9 @@ from flytekit.models.types import LiteralType, SimpleType, StructuredDatasetType, TypeStructure, UnionType try: - from typing import get_args as _get_args + from typing import Annotated, get_args, get_origin except ImportError: - try: - from typing_extensions import get_args as _get_args - except ImportError: - - def _get_args(t): - return t.__args__ + from typing_extensions import Annotated, get_args, get_origin T = typing.TypeVar("T") @@ -75,10 +65,13 @@ class TypeTransformer(typing.Generic[T]): Base transformer type that should be implemented for every python native type that can be handled by flytekit """ - def __init__(self, name: str, t: Type[T], enable_type_assertions: bool = True): + def __init__(self, name: str, t: Type[T], enable_type_assertions: bool = True, hash_overridable: bool = False): self._t = t self._name = name self._type_assertions_enabled = enable_type_assertions + # `hash_overridable` indicates that the literals produced by this type transformer can set their hashes if needed. + # See (link to documentation where this feature is explained). + self._hash_overridable = hash_overridable @property def name(self): @@ -98,6 +91,10 @@ def type_assertions_enabled(self) -> bool: """ return self._type_assertions_enabled + @property + def hash_overridable(self) -> bool: + return self._hash_overridable + def assert_type(self, t: Type[T], v: T): if not hasattr(t, "__origin__") and not isinstance(v, t): raise TypeTransformerFailedError(f"Type of Val '{v}' is not an instance of {t}") @@ -677,7 +674,25 @@ def to_literal(cls, ctx: FlyteContext, python_val: typing.Any, python_type: Type transformer = cls.get_transformer(python_type) if transformer.type_assertions_enabled: transformer.assert_type(python_type, python_val) + + # In case the value is an annotated type we inspect the annotations and look for hash-related annotations. + hash = None + if transformer.hash_overridable and get_origin(python_type) is Annotated: + # We are now dealing with one of two cases: + # 1. The annotated type is a `HashMethod`, which indicates that we should we should produce the hash using + # the method indicated in the annotation. + # 2. The annotated type is being used for a different purpose other than calculating hash values, in which case + # we should just continue. + for annotation in get_args(python_type)[1:]: + if not isinstance(annotation, HashMethod): + continue + hash = annotation.calculate(python_val) + break + lv = transformer.to_literal(ctx, python_val, python_type, expected) + + if hash is not None: + lv.hash = hash return lv @classmethod @@ -933,7 +948,7 @@ def __init__(self): def get_literal_type(self, t: Type[T]) -> Optional[LiteralType]: try: - trans = [(TypeEngine.get_transformer(x), x) for x in _get_args(t)] + trans = [(TypeEngine.get_transformer(x), x) for x in get_args(t)] variants = [_add_tag_to_type(t.get_literal_type(x), t.name) for (t, x) in trans] return _type_models.LiteralType(union_type=UnionType(variants)) except Exception as e: @@ -943,7 +958,7 @@ def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], exp found_res = False res = None res_type = None - for t in _get_args(python_type): + for t in get_args(python_type): try: trans = TypeEngine.get_transformer(t) @@ -973,7 +988,7 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: found_res = False res = None res_tag = None - for v in _get_args(expected_python_type): + for v in get_args(expected_python_type): try: trans = TypeEngine.get_transformer(v) if union_tag is not None: @@ -1080,6 +1095,7 @@ def to_literal( for k, v in python_val.items(): if type(k) != str: raise ValueError("Flyte MapType expects all keys to be strings") + # TODO: log a warning for Annotated objects that contain HashMethod k_type, v_type = self.get_dict_types(python_type) lit_map[k] = TypeEngine.to_literal(ctx, v, v_type, expected.map_value_type) return Literal(map=LiteralMap(literals=lit_map)) diff --git a/flytekit/models/literals.py b/flytekit/models/literals.py index c4ac6ca289..1bc69ae41b 100644 --- a/flytekit/models/literals.py +++ b/flytekit/models/literals.py @@ -851,7 +851,9 @@ def from_flyte_idl(cls, pb2_object): class Literal(_common.FlyteIdlEntity): - def __init__(self, scalar: Scalar = None, collection: LiteralCollection = None, map: LiteralMap = None): + def __init__( + self, scalar: Scalar = None, collection: LiteralCollection = None, map: LiteralMap = None, hash: str = None + ): """ This IDL message represents a literal value in the Flyte ecosystem. @@ -862,6 +864,7 @@ def __init__(self, scalar: Scalar = None, collection: LiteralCollection = None, self._scalar = scalar self._collection = collection self._map = map + self._hash = hash @property def scalar(self): @@ -895,6 +898,18 @@ def value(self): """ return self.scalar or self.collection or self.map + @property + def hash(self): + """ + If not None, this value holds a hash that represents the literal for caching purposes. + :rtype: str + """ + return self._hash + + @hash.setter + def hash(self, value): + self._hash = value + def to_flyte_idl(self): """ :rtype: flyteidl.core.literals_pb2.Literal @@ -903,6 +918,7 @@ def to_flyte_idl(self): scalar=self.scalar.to_flyte_idl() if self.scalar is not None else None, collection=self.collection.to_flyte_idl() if self.collection is not None else None, map=self.map.to_flyte_idl() if self.map is not None else None, + hash=self.hash, ) @classmethod @@ -919,4 +935,5 @@ def from_flyte_idl(cls, pb2_object): scalar=Scalar.from_flyte_idl(pb2_object.scalar) if pb2_object.HasField("scalar") else None, collection=collection, map=LiteralMap.from_flyte_idl(pb2_object.map) if pb2_object.HasField("map") else None, + hash=pb2_object.hash if pb2_object.hash else None, ) diff --git a/flytekit/types/schema/types_pandas.py b/flytekit/types/schema/types_pandas.py index 0edf024b08..fa0245236c 100644 --- a/flytekit/types/schema/types_pandas.py +++ b/flytekit/types/schema/types_pandas.py @@ -112,6 +112,9 @@ class PandasDataFrameTransformer(TypeTransformer[pandas.DataFrame]): def __init__(self): super().__init__("PandasDataFrame<->GenericSchema", pandas.DataFrame) self._parquet_engine = _PARQUETIO_ENGINES[sdk.PARQUET_ENGINE.get()] + # Pandas dataframes can have their hashes overriden to facilitate the case of caching pandas dataframes by + # value. + self._hash_overridable = True @staticmethod def _get_schema_type() -> SchemaType: diff --git a/flytekit/types/structured/structured_dataset.py b/flytekit/types/structured/structured_dataset.py index 7d9bdcb9ba..0b36a34302 100644 --- a/flytekit/types/structured/structured_dataset.py +++ b/flytekit/types/structured/structured_dataset.py @@ -381,6 +381,9 @@ def __init__(self): super().__init__("StructuredDataset Transformer", StructuredDataset) self._type_assertions_enabled = False + # Instances of StructuredDataset opt-in to the ability of being cached. + self._hash_overridable = True + @classmethod def register(cls, h: Handlers, default_for_type: Optional[bool] = True, override: Optional[bool] = False): """ diff --git a/plugins/flytekit-aws-athena/requirements.txt b/plugins/flytekit-aws-athena/requirements.txt index a1acae494f..bb210d8c13 100644 --- a/plugins/flytekit-aws-athena/requirements.txt +++ b/plugins/flytekit-aws-athena/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.9 +# This file is autogenerated by pip-compile with python 3.10 # To update, run: # # pip-compile requirements.in @@ -12,13 +12,15 @@ binaryornot==0.4.4 # via cookiecutter certifi==2021.10.8 # via requests +cffi==1.15.0 + # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.11 +charset-normalizer==2.0.12 # via requests checksumdir==1.2.0 # via flytekit -click==7.1.2 +click==8.0.4 # via # cookiecutter # flytekit @@ -26,7 +28,7 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.2.0 +croniter==1.3.4 # via flytekit cryptography==36.0.1 # via secretstorage @@ -42,16 +44,22 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.22.0 +flyteidl==0.23.0 # via flytekit -flytekit==0.30.0 +flytekit==0.30.3 # via flytekitplugins-athena -grpcio==1.43.0 +googleapis-common-protos==1.55.0 + # via flyteidl +grpcio==1.44.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.10.1 +importlib-metadata==4.11.2 # via keyring +jeepney==0.7.1 + # via + # keyring + # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -60,7 +68,7 @@ jinja2-time==0.2.0 # via cookiecutter keyring==23.5.0 # via flytekit -markupsafe==2.0.1 +markupsafe==2.1.0 # via jinja2 marshmallow==3.14.1 # via @@ -79,7 +87,7 @@ numpy==1.22.2 # via # pandas # pyarrow -pandas==1.4.0 +pandas==1.4.1 # via flytekit poyo==0.5.0 # via cookiecutter @@ -87,6 +95,10 @@ protobuf==3.19.4 # via # flyteidl # flytekit + # googleapis-common-protos + # protoc-gen-swagger +protoc-gen-swagger==0.1.0 + # via flyteidl py==1.11.0 # via retry pyarrow==6.0.1 @@ -101,7 +113,7 @@ python-dateutil==2.8.2 # pandas python-json-logger==2.0.2 # via flytekit -python-slugify==5.0.2 +python-slugify==6.1.1 # via cookiecutter pytimeparse==1.1.8 # via flytekit @@ -109,7 +121,7 @@ pytz==2021.3 # via # flytekit # pandas -regex==2022.1.18 +regex==2022.3.2 # via docker-image-py requests==2.27.1 # via @@ -120,6 +132,8 @@ responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit +secretstorage==3.3.1 + # via keyring six==1.16.0 # via # cookiecutter @@ -131,7 +145,7 @@ statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -typing-extensions==4.0.1 +typing-extensions==4.1.1 # via # flytekit # typing-inspect diff --git a/plugins/flytekit-aws-batch/requirements.txt b/plugins/flytekit-aws-batch/requirements.txt index 9181490234..80a3fcf51f 100644 --- a/plugins/flytekit-aws-batch/requirements.txt +++ b/plugins/flytekit-aws-batch/requirements.txt @@ -1,24 +1,26 @@ # -# This file is autogenerated by pip-compile with python 3.9 +# This file is autogenerated by pip-compile with python 3.10 # To update, run: # # pip-compile requirements.in # -e file:.#egg=flytekitplugins-awsbatch # via -r requirements.in -arrow==1.2.1 +arrow==1.2.2 # via jinja2-time binaryornot==0.4.4 # via cookiecutter certifi==2021.10.8 # via requests +cffi==1.15.0 + # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.10 +charset-normalizer==2.0.12 # via requests checksumdir==1.2.0 # via flytekit -click==7.1.2 +click==8.0.4 # via # cookiecutter # flytekit @@ -26,8 +28,10 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.2.0 +croniter==1.3.4 # via flytekit +cryptography==36.0.1 + # via secretstorage dataclasses-json==0.5.6 # via flytekit decorator==5.1.1 @@ -40,16 +44,22 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.21.23 +flyteidl==0.23.0 # via flytekit -flytekit==0.26.0 +flytekit==0.30.3 # via flytekitplugins-awsbatch -grpcio==1.43.0 +googleapis-common-protos==1.55.0 + # via flyteidl +grpcio==1.44.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.10.1 +importlib-metadata==4.11.2 # via keyring +jeepney==0.7.1 + # via + # keyring + # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -58,7 +68,7 @@ jinja2-time==0.2.0 # via cookiecutter keyring==23.5.0 # via flytekit -markupsafe==2.0.1 +markupsafe==2.1.0 # via jinja2 marshmallow==3.14.1 # via @@ -71,25 +81,31 @@ marshmallow-jsonschema==0.13.0 # via flytekit mypy-extensions==0.4.3 # via typing-inspect -natsort==8.0.2 +natsort==8.1.0 # via flytekit -numpy==1.22.1 +numpy==1.22.2 # via # pandas # pyarrow -pandas==1.3.5 +pandas==1.4.1 # via flytekit poyo==0.5.0 # via cookiecutter -protobuf==3.19.3 +protobuf==3.19.4 # via # flyteidl # flytekit + # googleapis-common-protos + # protoc-gen-swagger +protoc-gen-swagger==0.1.0 + # via flyteidl py==1.11.0 # via retry pyarrow==6.0.1 # via flytekit -python-dateutil==2.8.1 +pycparser==2.21 + # via cffi +python-dateutil==2.8.2 # via # arrow # croniter @@ -97,7 +113,7 @@ python-dateutil==2.8.1 # pandas python-json-logger==2.0.2 # via flytekit -python-slugify==5.0.2 +python-slugify==6.1.1 # via cookiecutter pytimeparse==1.1.8 # via flytekit @@ -105,32 +121,34 @@ pytz==2021.3 # via # flytekit # pandas -regex==2022.1.18 +regex==2022.3.2 # via docker-image-py requests==2.27.1 # via # cookiecutter # flytekit # responses -responses==0.17.0 +responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit +secretstorage==3.3.1 + # via keyring six==1.16.0 # via # cookiecutter - # flytekit # grpcio # python-dateutil - # responses sortedcontainers==2.4.0 # via flytekit statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -typing-extensions==4.0.1 - # via typing-inspect +typing-extensions==4.1.1 + # via + # flytekit + # typing-inspect typing-inspect==0.7.1 # via dataclasses-json urllib3==1.26.8 diff --git a/plugins/flytekit-aws-sagemaker/requirements.txt b/plugins/flytekit-aws-sagemaker/requirements.txt index 11396f213e..f4b02da229 100644 --- a/plugins/flytekit-aws-sagemaker/requirements.txt +++ b/plugins/flytekit-aws-sagemaker/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.9 +# This file is autogenerated by pip-compile with python 3.10 # To update, run: # # pip-compile requirements.in @@ -12,9 +12,9 @@ bcrypt==3.2.0 # via paramiko binaryornot==0.4.4 # via cookiecutter -boto3==1.20.50 +boto3==1.21.10 # via sagemaker-training -botocore==1.23.50 +botocore==1.24.10 # via # boto3 # s3transfer @@ -27,11 +27,11 @@ cffi==1.15.0 # pynacl chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.11 +charset-normalizer==2.0.12 # via requests checksumdir==1.2.0 # via flytekit -click==7.1.2 +click==8.0.4 # via # cookiecutter # flytekit @@ -39,7 +39,7 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.2.0 +croniter==1.3.4 # via flytekit cryptography==36.0.1 # via @@ -57,22 +57,28 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.22.0 +flyteidl==0.23.0 # via flytekit -flytekit==0.30.0 +flytekit==0.30.3 # via flytekitplugins-awssagemaker gevent==21.12.0 # via sagemaker-training +googleapis-common-protos==1.55.0 + # via flyteidl greenlet==1.1.2 # via gevent -grpcio==1.43.0 +grpcio==1.44.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.10.1 +importlib-metadata==4.11.2 # via keyring -inotify-simple==1.2.1 +inotify_simple==1.2.1 # via sagemaker-training +jeepney==0.7.1 + # via + # keyring + # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -85,7 +91,7 @@ jmespath==0.10.0 # botocore keyring==23.5.0 # via flytekit -markupsafe==2.0.1 +markupsafe==2.1.0 # via jinja2 marshmallow==3.14.1 # via @@ -106,7 +112,7 @@ numpy==1.22.2 # pyarrow # sagemaker-training # scipy -pandas==1.4.0 +pandas==1.4.1 # via flytekit paramiko==2.9.2 # via sagemaker-training @@ -116,7 +122,11 @@ protobuf==3.19.4 # via # flyteidl # flytekit + # googleapis-common-protos + # protoc-gen-swagger # sagemaker-training +protoc-gen-swagger==0.1.0 + # via flyteidl psutil==5.9.0 # via sagemaker-training py==1.11.0 @@ -136,7 +146,7 @@ python-dateutil==2.8.2 # pandas python-json-logger==2.0.2 # via flytekit -python-slugify==5.0.2 +python-slugify==6.1.1 # via cookiecutter pytimeparse==1.1.8 # via flytekit @@ -144,7 +154,7 @@ pytz==2021.3 # via # flytekit # pandas -regex==2022.1.18 +regex==2022.3.2 # via docker-image-py requests==2.27.1 # via @@ -157,12 +167,14 @@ retry==0.9.2 # via flytekit retrying==1.3.3 # via sagemaker-training -s3transfer==0.5.1 +s3transfer==0.5.2 # via boto3 sagemaker-training==3.9.2 # via flytekitplugins-awssagemaker scipy==1.8.0 # via sagemaker-training +secretstorage==3.3.1 + # via keyring six==1.16.0 # via # bcrypt @@ -177,7 +189,7 @@ statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -typing-extensions==4.0.1 +typing-extensions==4.1.1 # via # flytekit # typing-inspect @@ -199,9 +211,9 @@ wrapt==1.13.3 # flytekit zipp==3.7.0 # via importlib-metadata -zope-event==4.5.0 +zope.event==4.5.0 # via gevent -zope-interface==5.4.0 +zope.interface==5.4.0 # via gevent # The following packages are considered to be unsafe in a requirements file: diff --git a/plugins/flytekit-bigquery/requirements.txt b/plugins/flytekit-bigquery/requirements.txt index f0bc647453..aef0a38759 100644 --- a/plugins/flytekit-bigquery/requirements.txt +++ b/plugins/flytekit-bigquery/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.9 +# This file is autogenerated by pip-compile with python 3.10 # To update, run: # # pip-compile requirements.in @@ -18,11 +18,11 @@ cffi==1.15.0 # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.11 +charset-normalizer==2.0.12 # via requests checksumdir==1.2.0 # via flytekit -click==7.1.2 +click==8.0.4 # via # cookiecutter # flytekit @@ -30,7 +30,7 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.2.0 +croniter==1.3.4 # via flytekit cryptography==36.0.1 # via secretstorage @@ -46,9 +46,9 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.22.0 +flyteidl==0.23.0 # via flytekit -flytekit==0.30.0 +flytekit==0.30.3 # via flytekitplugins-bigquery google-api-core[grpc]==2.5.0 # via @@ -58,29 +58,30 @@ google-auth==2.6.0 # via # google-api-core # google-cloud-core -google-cloud-bigquery==2.32.0 +google-cloud-bigquery==2.34.0 # via flytekitplugins-bigquery google-cloud-core==2.2.2 # via google-cloud-bigquery google-crc32c==1.3.0 # via google-resumable-media -google-resumable-media==2.2.0 +google-resumable-media==2.3.0 # via google-cloud-bigquery -googleapis-common-protos==1.54.0 +googleapis-common-protos==1.55.0 # via + # flyteidl # google-api-core # grpcio-status -grpcio==1.43.0 +grpcio==1.44.0 # via # flytekit # google-api-core # google-cloud-bigquery # grpcio-status -grpcio-status==1.43.0 +grpcio-status==1.44.0 # via google-api-core idna==3.3 # via requests -importlib-metadata==4.10.1 +importlib-metadata==4.11.2 # via keyring jeepney==0.7.1 # via @@ -94,7 +95,7 @@ jinja2-time==0.2.0 # via cookiecutter keyring==23.5.0 # via flytekit -markupsafe==2.0.1 +markupsafe==2.1.0 # via jinja2 marshmallow==3.14.1 # via @@ -115,11 +116,11 @@ numpy==1.22.2 # pyarrow packaging==21.3 # via google-cloud-bigquery -pandas==1.4.0 +pandas==1.4.1 # via flytekit poyo==0.5.0 # via cookiecutter -proto-plus==1.20.0 +proto-plus==1.20.3 # via google-cloud-bigquery protobuf==3.19.4 # via @@ -130,6 +131,9 @@ protobuf==3.19.4 # googleapis-common-protos # grpcio-status # proto-plus + # protoc-gen-swagger +protoc-gen-swagger==0.1.0 + # via flyteidl py==1.11.0 # via retry pyarrow==6.0.1 @@ -153,7 +157,7 @@ python-dateutil==2.8.2 # pandas python-json-logger==2.0.2 # via flytekit -python-slugify==5.0.2 +python-slugify==6.1.1 # via cookiecutter pytimeparse==1.1.8 # via flytekit @@ -161,7 +165,7 @@ pytz==2021.3 # via # flytekit # pandas -regex==2022.1.18 +regex==2022.3.2 # via docker-image-py requests==2.27.1 # via @@ -190,7 +194,7 @@ statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -typing-extensions==4.0.1 +typing-extensions==4.1.1 # via # flytekit # typing-inspect diff --git a/plugins/flytekit-data-fsspec/requirements.txt b/plugins/flytekit-data-fsspec/requirements.txt index 4dc89fa500..b3f506d572 100644 --- a/plugins/flytekit-data-fsspec/requirements.txt +++ b/plugins/flytekit-data-fsspec/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.9 +# This file is autogenerated by pip-compile with python 3.10 # To update, run: # # pip-compile requirements.in @@ -10,17 +10,19 @@ arrow==1.2.2 # via jinja2-time binaryornot==0.4.4 # via cookiecutter -botocore==1.23.24 +botocore==1.24.10 # via flytekitplugins-data-fsspec certifi==2021.10.8 # via requests +cffi==1.15.0 + # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.11 +charset-normalizer==2.0.12 # via requests checksumdir==1.2.0 # via flytekit -click==7.1.2 +click==8.0.4 # via # cookiecutter # flytekit @@ -28,8 +30,10 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.2.0 +croniter==1.3.4 # via flytekit +cryptography==36.0.1 + # via secretstorage dataclasses-json==0.5.6 # via flytekit decorator==5.1.1 @@ -42,18 +46,24 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.22.0 +flyteidl==0.23.0 # via flytekit -flytekit==0.30.0 +flytekit==0.30.3 # via flytekitplugins-data-fsspec -fsspec==2022.1.0 +fsspec==2022.2.0 # via flytekitplugins-data-fsspec -grpcio==1.43.0 +googleapis-common-protos==1.55.0 + # via flyteidl +grpcio==1.44.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.10.1 +importlib-metadata==4.11.2 # via keyring +jeepney==0.7.1 + # via + # keyring + # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -64,7 +74,7 @@ jmespath==0.10.0 # via botocore keyring==23.5.0 # via flytekit -markupsafe==2.0.1 +markupsafe==2.1.0 # via jinja2 marshmallow==3.14.1 # via @@ -83,7 +93,7 @@ numpy==1.22.2 # via # pandas # pyarrow -pandas==1.4.0 +pandas==1.4.1 # via flytekit poyo==0.5.0 # via cookiecutter @@ -91,10 +101,16 @@ protobuf==3.19.4 # via # flyteidl # flytekit + # googleapis-common-protos + # protoc-gen-swagger +protoc-gen-swagger==0.1.0 + # via flyteidl py==1.11.0 # via retry pyarrow==6.0.1 # via flytekit +pycparser==2.21 + # via cffi python-dateutil==2.8.2 # via # arrow @@ -104,7 +120,7 @@ python-dateutil==2.8.2 # pandas python-json-logger==2.0.2 # via flytekit -python-slugify==5.0.2 +python-slugify==6.1.1 # via cookiecutter pytimeparse==1.1.8 # via flytekit @@ -112,7 +128,7 @@ pytz==2021.3 # via # flytekit # pandas -regex==2022.1.18 +regex==2022.3.2 # via docker-image-py requests==2.27.1 # via @@ -123,6 +139,8 @@ responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit +secretstorage==3.3.1 + # via keyring six==1.16.0 # via # cookiecutter @@ -134,7 +152,7 @@ statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -typing-extensions==4.0.1 +typing-extensions==4.1.1 # via # flytekit # typing-inspect diff --git a/plugins/flytekit-dolt/requirements.txt b/plugins/flytekit-dolt/requirements.txt index 01dfe2cb6b..fc9dc9f8b7 100644 --- a/plugins/flytekit-dolt/requirements.txt +++ b/plugins/flytekit-dolt/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.9 +# This file is autogenerated by pip-compile with python 3.10 # To update, run: # # pip-compile requirements.in @@ -12,13 +12,15 @@ binaryornot==0.4.4 # via cookiecutter certifi==2021.10.8 # via requests +cffi==1.15.0 + # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.11 +charset-normalizer==2.0.12 # via requests checksumdir==1.2.0 # via flytekit -click==7.1.2 +click==8.0.4 # via # cookiecutter # flytekit @@ -26,7 +28,7 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.2.0 +croniter==1.3.4 # via flytekit cryptography==36.0.1 # via secretstorage @@ -48,16 +50,22 @@ dolt-integrations==0.1.5 # via flytekitplugins-dolt doltcli==0.1.17 # via dolt-integrations -flyteidl==0.22.0 +flyteidl==0.23.0 # via flytekit -flytekit==0.30.0 +flytekit==0.30.3 # via flytekitplugins-dolt -grpcio==1.43.0 +googleapis-common-protos==1.55.0 + # via flyteidl +grpcio==1.44.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.10.1 +importlib-metadata==4.11.2 # via keyring +jeepney==0.7.1 + # via + # keyring + # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -66,7 +74,7 @@ jinja2-time==0.2.0 # via cookiecutter keyring==23.5.0 # via flytekit -markupsafe==2.0.1 +markupsafe==2.1.0 # via jinja2 marshmallow==3.14.1 # via @@ -85,7 +93,7 @@ numpy==1.22.2 # via # pandas # pyarrow -pandas==1.4.0 +pandas==1.4.1 # via # dolt-integrations # flytekit @@ -95,6 +103,10 @@ protobuf==3.19.4 # via # flyteidl # flytekit + # googleapis-common-protos + # protoc-gen-swagger +protoc-gen-swagger==0.1.0 + # via flyteidl py==1.11.0 # via retry pyarrow==6.0.1 @@ -109,7 +121,7 @@ python-dateutil==2.8.2 # pandas python-json-logger==2.0.2 # via flytekit -python-slugify==5.0.2 +python-slugify==6.1.1 # via cookiecutter pytimeparse==1.1.8 # via flytekit @@ -117,7 +129,7 @@ pytz==2021.3 # via # flytekit # pandas -regex==2022.1.18 +regex==2022.3.2 # via docker-image-py requests==2.27.1 # via @@ -128,6 +140,8 @@ responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit +secretstorage==3.3.1 + # via keyring six==1.16.0 # via # cookiecutter @@ -139,7 +153,7 @@ statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -typing-extensions==4.0.1 +typing-extensions==4.1.1 # via # flytekit # typing-inspect diff --git a/plugins/flytekit-greatexpectations/requirements.txt b/plugins/flytekit-greatexpectations/requirements.txt index 7aa0a867dd..31b3be4b48 100644 --- a/plugins/flytekit-greatexpectations/requirements.txt +++ b/plugins/flytekit-greatexpectations/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.9 +# This file is autogenerated by pip-compile with python 3.10 # To update, run: # # pip-compile requirements.in @@ -8,10 +8,6 @@ # via -r requirements.in altair==4.2.0 # via great-expectations -appnope==0.1.2 - # via - # ipykernel - # ipython argon2-cffi==21.3.0 # via notebook argon2-cffi-bindings==21.2.0 @@ -26,23 +22,22 @@ backcall==0.2.0 # via ipython binaryornot==0.4.4 # via cookiecutter -black==21.12b0 - # via ipython bleach==4.1.0 # via nbconvert certifi==2021.10.8 # via requests cffi==1.15.0 - # via argon2-cffi-bindings + # via + # argon2-cffi-bindings + # cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.10 +charset-normalizer==2.0.12 # via requests checksumdir==1.2.0 # via flytekit -click==7.1.2 +click==8.0.4 # via - # black # cookiecutter # flytekit # great-expectations @@ -50,8 +45,10 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.2.0 +croniter==1.3.4 # via flytekit +cryptography==36.0.1 + # via secretstorage dataclasses-json==0.5.6 # via flytekit debugpy==1.5.1 @@ -70,34 +67,36 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -entrypoints==0.3 +entrypoints==0.4 # via # altair # jupyter-client # nbconvert -executing==0.8.2 +executing==0.8.3 # via stack-data -flyteidl==0.21.24 +flyteidl==0.23.0 # via flytekit -flytekit==0.26.1 +flytekit==0.30.3 # via flytekitplugins-great-expectations -great-expectations==0.14.3 +googleapis-common-protos==1.55.0 + # via flyteidl +great-expectations==0.14.5 # via flytekitplugins-great-expectations greenlet==1.1.2 # via sqlalchemy -grpcio==1.43.0 +grpcio==1.44.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.10.1 +importlib-metadata==4.11.2 # via # great-expectations # keyring -ipykernel==6.7.0 +ipykernel==6.9.1 # via # ipywidgets # notebook -ipython==8.0.1 +ipython==8.1.0 # via # ipykernel # ipywidgets @@ -110,6 +109,10 @@ ipywidgets==7.6.5 # via great-expectations jedi==0.18.1 # via ipython +jeepney==0.7.1 + # via + # keyring + # secretstorage jinja2==3.0.3 # via # altair @@ -134,7 +137,7 @@ jupyter-client==7.1.2 # ipykernel # nbclient # notebook -jupyter-core==4.9.1 +jupyter-core==4.9.2 # via # jupyter-client # nbconvert @@ -146,7 +149,7 @@ jupyterlab-widgets==1.0.2 # via ipywidgets keyring==23.5.0 # via flytekit -markupsafe==2.0.1 +markupsafe==2.1.0 # via jinja2 marshmallow==3.14.1 # via @@ -166,14 +169,12 @@ mistune==0.8.4 # great-expectations # nbconvert mypy-extensions==0.4.3 - # via - # black - # typing-inspect -natsort==8.0.2 + # via typing-inspect +natsort==8.1.0 # via flytekit -nbclient==0.5.10 +nbclient==0.5.11 # via nbconvert -nbconvert==6.4.1 +nbconvert==6.4.2 # via notebook nbformat==5.1.3 # via @@ -189,7 +190,7 @@ nest-asyncio==1.5.4 # notebook notebook==6.4.8 # via widgetsnbextension -numpy==1.22.1 +numpy==1.22.2 # via # altair # great-expectations @@ -198,7 +199,7 @@ numpy==1.22.1 # scipy packaging==21.3 # via bleach -pandas==1.4.0 +pandas==1.4.1 # via # altair # flytekit @@ -207,24 +208,24 @@ pandocfilters==1.5.0 # via nbconvert parso==0.8.3 # via jedi -pathspec==0.9.0 - # via black pexpect==4.8.0 # via ipython pickleshare==0.7.5 # via ipython -platformdirs==2.4.1 - # via black poyo==0.5.0 # via cookiecutter prometheus-client==0.13.1 # via notebook -prompt-toolkit==3.0.26 +prompt-toolkit==3.0.28 # via ipython protobuf==3.19.4 # via # flyteidl # flytekit + # googleapis-common-protos + # protoc-gen-swagger +protoc-gen-swagger==0.1.0 + # via flyteidl ptyprocess==0.7.0 # via # pexpect @@ -248,7 +249,7 @@ pyparsing==2.4.7 # packaging pyrsistent==0.18.1 # via jsonschema -python-dateutil==2.8.1 +python-dateutil==2.8.2 # via # arrow # croniter @@ -258,7 +259,7 @@ python-dateutil==2.8.1 # pandas python-json-logger==2.0.2 # via flytekit -python-slugify==5.0.2 +python-slugify==6.1.1 # via cookiecutter pytimeparse==1.1.8 # via flytekit @@ -273,7 +274,7 @@ pyzmq==22.3.0 # via # jupyter-client # notebook -regex==2022.1.18 +regex==2022.3.2 # via docker-image-py requests==2.27.1 # via @@ -281,16 +282,16 @@ requests==2.27.1 # flytekit # great-expectations # responses -responses==0.17.0 +responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit -ruamel-yaml==0.17.17 +ruamel.yaml==0.17.17 # via great-expectations -ruamel-yaml-clib==0.2.6 - # via ruamel-yaml -scipy==1.7.3 +scipy==1.8.0 # via great-expectations +secretstorage==3.3.1 + # via keyring send2trash==1.8.0 # via notebook six==1.16.0 @@ -298,17 +299,15 @@ six==1.16.0 # asttokens # bleach # cookiecutter - # flytekit # grpcio # python-dateutil - # responses sortedcontainers==2.4.0 # via flytekit sqlalchemy==1.4.31 # via # -r requirements.in # flytekitplugins-great-expectations -stack-data==0.1.4 +stack-data==0.2.0 # via ipython statsd==3.3.0 # via flytekit @@ -316,12 +315,10 @@ termcolor==1.1.0 # via great-expectations terminado==0.13.1 # via notebook -testpath==0.5.0 +testpath==0.6.0 # via nbconvert text-unidecode==1.3 # via python-slugify -tomli==1.2.3 - # via black toolz==0.11.2 # via altair tornado==6.1 @@ -330,7 +327,7 @@ tornado==6.1 # jupyter-client # notebook # terminado -tqdm==4.62.3 +tqdm==4.63.0 # via great-expectations traitlets==5.1.1 # via @@ -344,9 +341,10 @@ traitlets==5.1.1 # nbconvert # nbformat # notebook -typing-extensions==4.0.1 +typing-extensions==4.1.1 # via - # black + # flytekit + # great-expectations # typing-inspect typing-inspect==0.7.1 # via dataclasses-json diff --git a/plugins/flytekit-hive/requirements.txt b/plugins/flytekit-hive/requirements.txt index 3ae057f9a0..69ed1e5f3e 100644 --- a/plugins/flytekit-hive/requirements.txt +++ b/plugins/flytekit-hive/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.9 +# This file is autogenerated by pip-compile with python 3.10 # To update, run: # # pip-compile requirements.in @@ -12,13 +12,15 @@ binaryornot==0.4.4 # via cookiecutter certifi==2021.10.8 # via requests +cffi==1.15.0 + # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.11 +charset-normalizer==2.0.12 # via requests checksumdir==1.2.0 # via flytekit -click==7.1.2 +click==8.0.4 # via # cookiecutter # flytekit @@ -26,7 +28,7 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.2.0 +croniter==1.3.4 # via flytekit cryptography==36.0.1 # via secretstorage @@ -42,16 +44,22 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.22.0 +flyteidl==0.23.0 # via flytekit -flytekit==0.30.0 +flytekit==0.30.3 # via flytekitplugins-hive -grpcio==1.43.0 +googleapis-common-protos==1.55.0 + # via flyteidl +grpcio==1.44.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.10.1 +importlib-metadata==4.11.2 # via keyring +jeepney==0.7.1 + # via + # keyring + # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -60,7 +68,7 @@ jinja2-time==0.2.0 # via cookiecutter keyring==23.5.0 # via flytekit -markupsafe==2.0.1 +markupsafe==2.1.0 # via jinja2 marshmallow==3.14.1 # via @@ -79,7 +87,7 @@ numpy==1.22.2 # via # pandas # pyarrow -pandas==1.4.0 +pandas==1.4.1 # via flytekit poyo==0.5.0 # via cookiecutter @@ -87,6 +95,10 @@ protobuf==3.19.4 # via # flyteidl # flytekit + # googleapis-common-protos + # protoc-gen-swagger +protoc-gen-swagger==0.1.0 + # via flyteidl py==1.11.0 # via retry pyarrow==6.0.1 @@ -101,7 +113,7 @@ python-dateutil==2.8.2 # pandas python-json-logger==2.0.2 # via flytekit -python-slugify==5.0.2 +python-slugify==6.1.1 # via cookiecutter pytimeparse==1.1.8 # via flytekit @@ -109,7 +121,7 @@ pytz==2021.3 # via # flytekit # pandas -regex==2022.1.18 +regex==2022.3.2 # via docker-image-py requests==2.27.1 # via @@ -120,6 +132,8 @@ responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit +secretstorage==3.3.1 + # via keyring six==1.16.0 # via # cookiecutter @@ -131,7 +145,7 @@ statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -typing-extensions==4.0.1 +typing-extensions==4.1.1 # via # flytekit # typing-inspect diff --git a/plugins/flytekit-k8s-pod/requirements.txt b/plugins/flytekit-k8s-pod/requirements.txt index d60d9c46b6..df398913af 100644 --- a/plugins/flytekit-k8s-pod/requirements.txt +++ b/plugins/flytekit-k8s-pod/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.9 +# This file is autogenerated by pip-compile with python 3.10 # To update, run: # # pip-compile requirements.in @@ -16,13 +16,15 @@ certifi==2021.10.8 # via # kubernetes # requests +cffi==1.15.0 + # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.11 +charset-normalizer==2.0.12 # via requests checksumdir==1.2.0 # via flytekit -click==7.1.2 +click==8.0.4 # via # cookiecutter # flytekit @@ -30,7 +32,7 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.2.0 +croniter==1.3.4 # via flytekit cryptography==36.0.1 # via secretstorage @@ -46,18 +48,24 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.22.0 +flyteidl==0.23.0 # via flytekit -flytekit==0.30.0 +flytekit==0.30.3 # via flytekitplugins-pod google-auth==2.6.0 # via kubernetes -grpcio==1.43.0 +googleapis-common-protos==1.55.0 + # via flyteidl +grpcio==1.44.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.10.1 +importlib-metadata==4.11.2 # via keyring +jeepney==0.7.1 + # via + # keyring + # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -66,9 +74,9 @@ jinja2-time==0.2.0 # via cookiecutter keyring==23.5.0 # via flytekit -kubernetes==21.7.0 +kubernetes==23.3.0 # via flytekitplugins-pod -markupsafe==2.0.1 +markupsafe==2.1.0 # via jinja2 marshmallow==3.14.1 # via @@ -89,7 +97,7 @@ numpy==1.22.2 # pyarrow oauthlib==3.2.0 # via requests-oauthlib -pandas==1.4.0 +pandas==1.4.1 # via flytekit poyo==0.5.0 # via cookiecutter @@ -97,6 +105,10 @@ protobuf==3.19.4 # via # flyteidl # flytekit + # googleapis-common-protos + # protoc-gen-swagger +protoc-gen-swagger==0.1.0 + # via flyteidl py==1.11.0 # via retry pyarrow==6.0.1 @@ -118,7 +130,7 @@ python-dateutil==2.8.2 # pandas python-json-logger==2.0.2 # via flytekit -python-slugify==5.0.2 +python-slugify==6.1.1 # via cookiecutter pytimeparse==1.1.8 # via flytekit @@ -128,7 +140,7 @@ pytz==2021.3 # pandas pyyaml==6.0 # via kubernetes -regex==2022.1.18 +regex==2022.3.2 # via docker-image-py requests==2.27.1 # via @@ -145,6 +157,8 @@ retry==0.9.2 # via flytekit rsa==4.8 # via google-auth +secretstorage==3.3.1 + # via keyring six==1.16.0 # via # cookiecutter @@ -158,7 +172,7 @@ statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -typing-extensions==4.0.1 +typing-extensions==4.1.1 # via # flytekit # typing-inspect @@ -170,7 +184,7 @@ urllib3==1.26.8 # kubernetes # requests # responses -websocket-client==1.2.3 +websocket-client==1.3.1 # via kubernetes wheel==0.37.1 # via flytekit diff --git a/plugins/flytekit-kf-mpi/requirements.txt b/plugins/flytekit-kf-mpi/requirements.txt index c13f9f3b74..82a2317f2f 100644 --- a/plugins/flytekit-kf-mpi/requirements.txt +++ b/plugins/flytekit-kf-mpi/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.9 +# This file is autogenerated by pip-compile with python 3.10 # To update, run: # # pip-compile requirements.in @@ -12,13 +12,15 @@ binaryornot==0.4.4 # via cookiecutter certifi==2021.10.8 # via requests +cffi==1.15.0 + # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.11 +charset-normalizer==2.0.12 # via requests checksumdir==1.2.0 # via flytekit -click==7.1.2 +click==8.0.4 # via # cookiecutter # flytekit @@ -26,7 +28,7 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.2.0 +croniter==1.3.4 # via flytekit cryptography==36.0.1 # via secretstorage @@ -42,18 +44,24 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.22.0 +flyteidl==0.23.0 # via # flytekit # flytekitplugins-kfmpi -flytekit==0.30.0 +flytekit==0.30.3 # via flytekitplugins-kfmpi -grpcio==1.43.0 +googleapis-common-protos==1.55.0 + # via flyteidl +grpcio==1.44.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.10.1 +importlib-metadata==4.11.2 # via keyring +jeepney==0.7.1 + # via + # keyring + # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -62,7 +70,7 @@ jinja2-time==0.2.0 # via cookiecutter keyring==23.5.0 # via flytekit -markupsafe==2.0.1 +markupsafe==2.1.0 # via jinja2 marshmallow==3.14.1 # via @@ -81,7 +89,7 @@ numpy==1.22.2 # via # pandas # pyarrow -pandas==1.4.0 +pandas==1.4.1 # via flytekit poyo==0.5.0 # via cookiecutter @@ -89,6 +97,10 @@ protobuf==3.19.4 # via # flyteidl # flytekit + # googleapis-common-protos + # protoc-gen-swagger +protoc-gen-swagger==0.1.0 + # via flyteidl py==1.11.0 # via retry pyarrow==6.0.1 @@ -103,7 +115,7 @@ python-dateutil==2.8.2 # pandas python-json-logger==2.0.2 # via flytekit -python-slugify==5.0.2 +python-slugify==6.1.1 # via cookiecutter pytimeparse==1.1.8 # via flytekit @@ -111,7 +123,7 @@ pytz==2021.3 # via # flytekit # pandas -regex==2022.1.18 +regex==2022.3.2 # via docker-image-py requests==2.27.1 # via @@ -122,6 +134,8 @@ responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit +secretstorage==3.3.1 + # via keyring six==1.16.0 # via # cookiecutter @@ -133,7 +147,7 @@ statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -typing-extensions==4.0.1 +typing-extensions==4.1.1 # via # flytekit # typing-inspect diff --git a/plugins/flytekit-kf-pytorch/requirements.txt b/plugins/flytekit-kf-pytorch/requirements.txt index 1696646c87..0a53a62333 100644 --- a/plugins/flytekit-kf-pytorch/requirements.txt +++ b/plugins/flytekit-kf-pytorch/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.9 +# This file is autogenerated by pip-compile with python 3.10 # To update, run: # # pip-compile requirements.in @@ -12,13 +12,15 @@ binaryornot==0.4.4 # via cookiecutter certifi==2021.10.8 # via requests +cffi==1.15.0 + # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.11 +charset-normalizer==2.0.12 # via requests checksumdir==1.2.0 # via flytekit -click==7.1.2 +click==8.0.4 # via # cookiecutter # flytekit @@ -26,7 +28,7 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.2.0 +croniter==1.3.4 # via flytekit cryptography==36.0.1 # via secretstorage @@ -42,16 +44,22 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.22.0 +flyteidl==0.23.0 # via flytekit -flytekit==0.30.0 +flytekit==0.30.3 # via flytekitplugins-kfpytorch -grpcio==1.43.0 +googleapis-common-protos==1.55.0 + # via flyteidl +grpcio==1.44.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.10.1 +importlib-metadata==4.11.2 # via keyring +jeepney==0.7.1 + # via + # keyring + # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -60,7 +68,7 @@ jinja2-time==0.2.0 # via cookiecutter keyring==23.5.0 # via flytekit -markupsafe==2.0.1 +markupsafe==2.1.0 # via jinja2 marshmallow==3.14.1 # via @@ -79,7 +87,7 @@ numpy==1.22.2 # via # pandas # pyarrow -pandas==1.4.0 +pandas==1.4.1 # via flytekit poyo==0.5.0 # via cookiecutter @@ -87,6 +95,10 @@ protobuf==3.19.4 # via # flyteidl # flytekit + # googleapis-common-protos + # protoc-gen-swagger +protoc-gen-swagger==0.1.0 + # via flyteidl py==1.11.0 # via retry pyarrow==6.0.1 @@ -101,7 +113,7 @@ python-dateutil==2.8.2 # pandas python-json-logger==2.0.2 # via flytekit -python-slugify==5.0.2 +python-slugify==6.1.1 # via cookiecutter pytimeparse==1.1.8 # via flytekit @@ -109,7 +121,7 @@ pytz==2021.3 # via # flytekit # pandas -regex==2022.1.18 +regex==2022.3.2 # via docker-image-py requests==2.27.1 # via @@ -120,6 +132,8 @@ responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit +secretstorage==3.3.1 + # via keyring six==1.16.0 # via # cookiecutter @@ -131,7 +145,7 @@ statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -typing-extensions==4.0.1 +typing-extensions==4.1.1 # via # flytekit # typing-inspect diff --git a/plugins/flytekit-kf-tensorflow/requirements.txt b/plugins/flytekit-kf-tensorflow/requirements.txt index d3b6253664..b782c39630 100644 --- a/plugins/flytekit-kf-tensorflow/requirements.txt +++ b/plugins/flytekit-kf-tensorflow/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.9 +# This file is autogenerated by pip-compile with python 3.10 # To update, run: # # pip-compile requirements.in @@ -12,13 +12,15 @@ binaryornot==0.4.4 # via cookiecutter certifi==2021.10.8 # via requests +cffi==1.15.0 + # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.11 +charset-normalizer==2.0.12 # via requests checksumdir==1.2.0 # via flytekit -click==7.1.2 +click==8.0.4 # via # cookiecutter # flytekit @@ -26,7 +28,7 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.2.0 +croniter==1.3.4 # via flytekit cryptography==36.0.1 # via secretstorage @@ -42,16 +44,22 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.22.0 +flyteidl==0.23.0 # via flytekit -flytekit==0.30.0 +flytekit==0.30.3 # via flytekitplugins-kftensorflow -grpcio==1.43.0 +googleapis-common-protos==1.55.0 + # via flyteidl +grpcio==1.44.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.10.1 +importlib-metadata==4.11.2 # via keyring +jeepney==0.7.1 + # via + # keyring + # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -60,7 +68,7 @@ jinja2-time==0.2.0 # via cookiecutter keyring==23.5.0 # via flytekit -markupsafe==2.0.1 +markupsafe==2.1.0 # via jinja2 marshmallow==3.14.1 # via @@ -79,7 +87,7 @@ numpy==1.22.2 # via # pandas # pyarrow -pandas==1.4.0 +pandas==1.4.1 # via flytekit poyo==0.5.0 # via cookiecutter @@ -87,6 +95,10 @@ protobuf==3.19.4 # via # flyteidl # flytekit + # googleapis-common-protos + # protoc-gen-swagger +protoc-gen-swagger==0.1.0 + # via flyteidl py==1.11.0 # via retry pyarrow==6.0.1 @@ -101,7 +113,7 @@ python-dateutil==2.8.2 # pandas python-json-logger==2.0.2 # via flytekit -python-slugify==5.0.2 +python-slugify==6.1.1 # via cookiecutter pytimeparse==1.1.8 # via flytekit @@ -109,7 +121,7 @@ pytz==2021.3 # via # flytekit # pandas -regex==2022.1.18 +regex==2022.3.2 # via docker-image-py requests==2.27.1 # via @@ -120,6 +132,8 @@ responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit +secretstorage==3.3.1 + # via keyring six==1.16.0 # via # cookiecutter @@ -131,7 +145,7 @@ statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -typing-extensions==4.0.1 +typing-extensions==4.1.1 # via # flytekit # typing-inspect diff --git a/plugins/flytekit-pandera/requirements.txt b/plugins/flytekit-pandera/requirements.txt index f0b65cec17..7b72a0b4c0 100644 --- a/plugins/flytekit-pandera/requirements.txt +++ b/plugins/flytekit-pandera/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.9 +# This file is autogenerated by pip-compile with python 3.10 # To update, run: # # pip-compile requirements.in @@ -12,13 +12,15 @@ binaryornot==0.4.4 # via cookiecutter certifi==2021.10.8 # via requests +cffi==1.15.0 + # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.11 +charset-normalizer==2.0.12 # via requests checksumdir==1.2.0 # via flytekit -click==7.1.2 +click==8.0.4 # via # cookiecutter # flytekit @@ -26,7 +28,7 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.2.0 +croniter==1.3.4 # via flytekit cryptography==36.0.1 # via secretstorage @@ -42,16 +44,22 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.22.0 +flyteidl==0.23.0 # via flytekit -flytekit==0.30.0 +flytekit==0.30.3 # via flytekitplugins-pandera -grpcio==1.43.0 +googleapis-common-protos==1.55.0 + # via flyteidl +grpcio==1.44.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.10.1 +importlib-metadata==4.11.2 # via keyring +jeepney==0.7.1 + # via + # keyring + # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -60,7 +68,7 @@ jinja2-time==0.2.0 # via cookiecutter keyring==23.5.0 # via flytekit -markupsafe==2.0.1 +markupsafe==2.1.0 # via jinja2 marshmallow==3.14.1 # via @@ -82,11 +90,11 @@ numpy==1.22.2 # pyarrow packaging==21.3 # via pandera -pandas==1.4.0 +pandas==1.4.1 # via # flytekit # pandera -pandera==0.8.1 +pandera==0.9.0 # via flytekitplugins-pandera poyo==0.5.0 # via cookiecutter @@ -94,6 +102,10 @@ protobuf==3.19.4 # via # flyteidl # flytekit + # googleapis-common-protos + # protoc-gen-swagger +protoc-gen-swagger==0.1.0 + # via flyteidl py==1.11.0 # via retry pyarrow==6.0.1 @@ -102,6 +114,8 @@ pyarrow==6.0.1 # pandera pycparser==2.21 # via cffi +pydantic==1.9.0 + # via pandera pyparsing==3.0.7 # via packaging python-dateutil==2.8.2 @@ -112,7 +126,7 @@ python-dateutil==2.8.2 # pandas python-json-logger==2.0.2 # via flytekit -python-slugify==5.0.2 +python-slugify==6.1.1 # via cookiecutter pytimeparse==1.1.8 # via flytekit @@ -120,7 +134,7 @@ pytz==2021.3 # via # flytekit # pandas -regex==2022.1.18 +regex==2022.3.2 # via docker-image-py requests==2.27.1 # via @@ -131,6 +145,8 @@ responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit +secretstorage==3.3.1 + # via keyring six==1.16.0 # via # cookiecutter @@ -142,9 +158,10 @@ statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -typing-extensions==4.0.1 +typing-extensions==4.1.1 # via # flytekit + # pydantic # typing-inspect typing-inspect==0.7.1 # via diff --git a/plugins/flytekit-papermill/dev-requirements.in b/plugins/flytekit-papermill/dev-requirements.in index 32404e0a2d..7a47470868 100644 --- a/plugins/flytekit-papermill/dev-requirements.in +++ b/plugins/flytekit-papermill/dev-requirements.in @@ -1,2 +1,3 @@ -git+https://github.com/flyteorg/flytekit@add-sd-make-class-methods#egg=flytekitplugins-spark&subdirectory=plugins/flytekit-spark +flyteidl>=0.23.0 +git+https://github.com/flyteorg/flytekit@offloaded-objects-caching#egg=flytekitplugins-spark&subdirectory=plugins/flytekit-spark # vcs+protocol://repo_url/#egg=pkg&subdirectory=flyte diff --git a/plugins/flytekit-papermill/dev-requirements.txt b/plugins/flytekit-papermill/dev-requirements.txt index ef30545338..f777991cc3 100644 --- a/plugins/flytekit-papermill/dev-requirements.txt +++ b/plugins/flytekit-papermill/dev-requirements.txt @@ -10,6 +10,8 @@ binaryornot==0.4.4 # via cookiecutter certifi==2021.10.8 # via requests +cffi==1.15.0 + # via cryptography chardet==4.0.0 # via binaryornot charset-normalizer==2.0.10 @@ -26,6 +28,8 @@ cookiecutter==1.7.3 # via flytekit croniter==1.2.0 # via flytekit +cryptography==36.0.1 + # via secretstorage dataclasses-json==0.5.6 # via flytekit decorator==5.1.1 @@ -38,18 +42,26 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.22.0 - # via flytekit +flyteidl==0.23.0 + # via + # -r dev-requirements.in + # flytekit flytekit==0.30.0 # via flytekitplugins-spark -flytekitplugins-spark @ git+https://github.com/flyteorg/flytekit@add-sd-make-class-methods#subdirectory=plugins/flytekit-spark +flytekitplugins-spark @ git+https://github.com/flyteorg/flytekit@offloaded-objects-caching#subdirectory=plugins/flytekit-spark # via -r dev-requirements.in +googleapis-common-protos==1.55.0 + # via flyteidl grpcio==1.43.0 # via flytekit idna==3.3 # via requests importlib-metadata==4.10.1 # via keyring +jeepney==0.7.1 + # via + # keyring + # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -85,12 +97,18 @@ protobuf==3.19.3 # via # flyteidl # flytekit + # googleapis-common-protos + # protoc-gen-swagger +protoc-gen-swagger==0.1.0 + # via flyteidl py==1.11.0 # via retry py4j==0.10.9.3 # via pyspark pyarrow==6.0.1 # via flytekit +pycparser==2.21 + # via cffi pyspark==3.2.1 # via flytekitplugins-spark python-dateutil==2.8.1 @@ -120,6 +138,8 @@ responses==0.17.0 # via flytekit retry==0.9.2 # via flytekit +secretstorage==3.3.1 + # via keyring six==1.16.0 # via # cookiecutter diff --git a/plugins/flytekit-snowflake/requirements.txt b/plugins/flytekit-snowflake/requirements.txt index 53951100b5..2d34ec966f 100644 --- a/plugins/flytekit-snowflake/requirements.txt +++ b/plugins/flytekit-snowflake/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.9 +# This file is autogenerated by pip-compile with python 3.10 # To update, run: # # pip-compile requirements.in @@ -12,13 +12,15 @@ binaryornot==0.4.4 # via cookiecutter certifi==2021.10.8 # via requests +cffi==1.15.0 + # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.11 +charset-normalizer==2.0.12 # via requests checksumdir==1.2.0 # via flytekit -click==7.1.2 +click==8.0.4 # via # cookiecutter # flytekit @@ -26,7 +28,7 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.2.0 +croniter==1.3.4 # via flytekit cryptography==36.0.1 # via secretstorage @@ -42,16 +44,22 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.22.0 +flyteidl==0.23.0 # via flytekit -flytekit==0.30.0 +flytekit==0.30.3 # via flytekitplugins-snowflake -grpcio==1.43.0 +googleapis-common-protos==1.55.0 + # via flyteidl +grpcio==1.44.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.10.1 +importlib-metadata==4.11.2 # via keyring +jeepney==0.7.1 + # via + # keyring + # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -60,7 +68,7 @@ jinja2-time==0.2.0 # via cookiecutter keyring==23.5.0 # via flytekit -markupsafe==2.0.1 +markupsafe==2.1.0 # via jinja2 marshmallow==3.14.1 # via @@ -79,7 +87,7 @@ numpy==1.22.2 # via # pandas # pyarrow -pandas==1.4.0 +pandas==1.4.1 # via flytekit poyo==0.5.0 # via cookiecutter @@ -87,6 +95,10 @@ protobuf==3.19.4 # via # flyteidl # flytekit + # googleapis-common-protos + # protoc-gen-swagger +protoc-gen-swagger==0.1.0 + # via flyteidl py==1.11.0 # via retry pyarrow==6.0.1 @@ -101,7 +113,7 @@ python-dateutil==2.8.2 # pandas python-json-logger==2.0.2 # via flytekit -python-slugify==5.0.2 +python-slugify==6.1.1 # via cookiecutter pytimeparse==1.1.8 # via flytekit @@ -109,7 +121,7 @@ pytz==2021.3 # via # flytekit # pandas -regex==2022.1.18 +regex==2022.3.2 # via docker-image-py requests==2.27.1 # via @@ -120,6 +132,8 @@ responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit +secretstorage==3.3.1 + # via keyring six==1.16.0 # via # cookiecutter @@ -131,7 +145,7 @@ statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -typing-extensions==4.0.1 +typing-extensions==4.1.1 # via # flytekit # typing-inspect diff --git a/plugins/flytekit-spark/requirements.txt b/plugins/flytekit-spark/requirements.txt index 75d1e32d53..6e7c6e59f0 100644 --- a/plugins/flytekit-spark/requirements.txt +++ b/plugins/flytekit-spark/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.9 +# This file is autogenerated by pip-compile with python 3.10 # To update, run: # # pip-compile requirements.in @@ -12,13 +12,15 @@ binaryornot==0.4.4 # via cookiecutter certifi==2021.10.8 # via requests +cffi==1.15.0 + # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.11 +charset-normalizer==2.0.12 # via requests checksumdir==1.2.0 # via flytekit -click==7.1.2 +click==8.0.4 # via # cookiecutter # flytekit @@ -26,7 +28,7 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.2.0 +croniter==1.3.4 # via flytekit cryptography==36.0.1 # via secretstorage @@ -42,16 +44,22 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.22.0 +flyteidl==0.23.0 # via flytekit -flytekit==0.30.0 +flytekit==0.30.3 # via flytekitplugins-spark -grpcio==1.43.0 +googleapis-common-protos==1.55.0 + # via flyteidl +grpcio==1.44.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.10.1 +importlib-metadata==4.11.2 # via keyring +jeepney==0.7.1 + # via + # keyring + # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -60,7 +68,7 @@ jinja2-time==0.2.0 # via cookiecutter keyring==23.5.0 # via flytekit -markupsafe==2.0.1 +markupsafe==2.1.0 # via jinja2 marshmallow==3.14.1 # via @@ -79,7 +87,7 @@ numpy==1.22.2 # via # pandas # pyarrow -pandas==1.4.0 +pandas==1.4.1 # via flytekit poyo==0.5.0 # via cookiecutter @@ -87,6 +95,10 @@ protobuf==3.19.4 # via # flyteidl # flytekit + # googleapis-common-protos + # protoc-gen-swagger +protoc-gen-swagger==0.1.0 + # via flyteidl py==1.11.0 # via retry py4j==0.10.9.3 @@ -105,7 +117,7 @@ python-dateutil==2.8.2 # pandas python-json-logger==2.0.2 # via flytekit -python-slugify==5.0.2 +python-slugify==6.1.1 # via cookiecutter pytimeparse==1.1.8 # via flytekit @@ -113,7 +125,7 @@ pytz==2021.3 # via # flytekit # pandas -regex==2022.1.18 +regex==2022.3.2 # via docker-image-py requests==2.27.1 # via @@ -124,6 +136,8 @@ responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit +secretstorage==3.3.1 + # via keyring six==1.16.0 # via # cookiecutter @@ -135,7 +149,7 @@ statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -typing-extensions==4.0.1 +typing-extensions==4.1.1 # via # flytekit # typing-inspect diff --git a/plugins/flytekit-sqlalchemy/requirements.txt b/plugins/flytekit-sqlalchemy/requirements.txt index 93dcc2790c..e617920549 100644 --- a/plugins/flytekit-sqlalchemy/requirements.txt +++ b/plugins/flytekit-sqlalchemy/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.9 +# This file is autogenerated by pip-compile with python 3.10 # To update, run: # # pip-compile requirements.in @@ -12,13 +12,15 @@ binaryornot==0.4.4 # via cookiecutter certifi==2021.10.8 # via requests +cffi==1.15.0 + # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.11 +charset-normalizer==2.0.12 # via requests checksumdir==1.2.0 # via flytekit -click==7.1.2 +click==8.0.4 # via # cookiecutter # flytekit @@ -26,7 +28,7 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.2.0 +croniter==1.3.4 # via flytekit cryptography==36.0.1 # via secretstorage @@ -42,18 +44,24 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.22.0 +flyteidl==0.23.0 # via flytekit -flytekit==0.30.0 +flytekit==0.30.3 # via flytekitplugins-sqlalchemy +googleapis-common-protos==1.55.0 + # via flyteidl greenlet==1.1.2 # via sqlalchemy -grpcio==1.43.0 +grpcio==1.44.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.10.1 +importlib-metadata==4.11.2 # via keyring +jeepney==0.7.1 + # via + # keyring + # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -62,7 +70,7 @@ jinja2-time==0.2.0 # via cookiecutter keyring==23.5.0 # via flytekit -markupsafe==2.0.1 +markupsafe==2.1.0 # via jinja2 marshmallow==3.14.1 # via @@ -81,7 +89,7 @@ numpy==1.22.2 # via # pandas # pyarrow -pandas==1.4.0 +pandas==1.4.1 # via flytekit poyo==0.5.0 # via cookiecutter @@ -89,6 +97,10 @@ protobuf==3.19.4 # via # flyteidl # flytekit + # googleapis-common-protos + # protoc-gen-swagger +protoc-gen-swagger==0.1.0 + # via flyteidl py==1.11.0 # via retry pyarrow==6.0.1 @@ -103,7 +115,7 @@ python-dateutil==2.8.2 # pandas python-json-logger==2.0.2 # via flytekit -python-slugify==5.0.2 +python-slugify==6.1.1 # via cookiecutter pytimeparse==1.1.8 # via flytekit @@ -111,7 +123,7 @@ pytz==2021.3 # via # flytekit # pandas -regex==2022.1.18 +regex==2022.3.2 # via docker-image-py requests==2.27.1 # via @@ -122,6 +134,8 @@ responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit +secretstorage==3.3.1 + # via keyring six==1.16.0 # via # cookiecutter @@ -135,7 +149,7 @@ statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -typing-extensions==4.0.1 +typing-extensions==4.1.1 # via # flytekit # typing-inspect diff --git a/requirements-spark2.txt b/requirements-spark2.txt index 3a007799e6..bd072c15d6 100644 --- a/requirements-spark2.txt +++ b/requirements-spark2.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.8 +# This file is autogenerated by pip-compile with python 3.9 # To update, run: # # make requirements-spark2.txt @@ -20,6 +20,8 @@ binaryornot==0.4.4 # via cookiecutter certifi==2021.10.8 # via requests +cffi==1.15.0 + # via cryptography chardet==4.0.0 # via binaryornot charset-normalizer==2.0.12 @@ -36,6 +38,8 @@ cookiecutter==1.7.3 # via flytekit croniter==1.3.4 # via flytekit +cryptography==36.0.1 + # via secretstorage dataclasses-json==0.5.6 # via flytekit decorator==5.1.1 @@ -48,9 +52,9 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.22.3 +flyteidl==0.23.0 # via flytekit -googleapis-common-protos==1.54.0 +googleapis-common-protos==1.55.0 # via flyteidl grpcio==1.44.0 # via flytekit @@ -58,6 +62,10 @@ idna==3.3 # via requests importlib-metadata==4.11.2 # via keyring +jeepney==0.7.1 + # via + # keyring + # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -106,6 +114,8 @@ py==1.11.0 # via retry pyarrow==6.0.1 # via flytekit +pycparser==2.21 + # via cffi pyrsistent==0.18.1 # via jsonschema python-dateutil==2.8.2 @@ -137,6 +147,8 @@ responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit +secretstorage==3.3.1 + # via keyring six==1.16.0 # via # cookiecutter diff --git a/requirements.txt b/requirements.txt index 4c51a9362a..0161dded53 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.8 +# This file is autogenerated by pip-compile with python 3.9 # To update, run: # # make requirements.txt @@ -18,6 +18,8 @@ binaryornot==0.4.4 # via cookiecutter certifi==2021.10.8 # via requests +cffi==1.15.0 + # via cryptography chardet==4.0.0 # via binaryornot charset-normalizer==2.0.12 @@ -34,6 +36,8 @@ cookiecutter==1.7.3 # via flytekit croniter==1.3.4 # via flytekit +cryptography==36.0.1 + # via secretstorage dataclasses-json==0.5.6 # via flytekit decorator==5.1.1 @@ -46,9 +50,9 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.22.3 +flyteidl==0.23.0 # via flytekit -googleapis-common-protos==1.54.0 +googleapis-common-protos==1.55.0 # via flyteidl grpcio==1.44.0 # via flytekit @@ -56,6 +60,10 @@ idna==3.3 # via requests importlib-metadata==4.11.2 # via keyring +jeepney==0.7.1 + # via + # keyring + # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -104,6 +112,8 @@ py==1.11.0 # via retry pyarrow==6.0.1 # via flytekit +pycparser==2.21 + # via cffi pyrsistent==0.18.1 # via jsonschema python-dateutil==2.8.2 @@ -115,10 +125,14 @@ python-dateutil==2.8.2 python-json-logger==2.0.2 # via flytekit <<<<<<< HEAD +<<<<<<< HEAD python-slugify==6.1.1 ======= python-slugify==6.1.0 >>>>>>> 6a1ceea5 (Bump idl (#862)) +======= +python-slugify==6.1.1 +>>>>>>> 0da523c7 (Caching of offloaded objects (#762)) # via cookiecutter pytimeparse==1.1.8 # via flytekit @@ -139,6 +153,8 @@ responses==0.18.0 # via flytekit retry==0.9.2 # via flytekit +secretstorage==3.3.1 + # via keyring six==1.16.0 # via # cookiecutter diff --git a/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt b/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt index 22200b321a..fb365ed51d 100644 --- a/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt +++ b/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.8 +# This file is autogenerated by pip-compile with python 3.9 # To update, run: # # make tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt @@ -119,6 +119,8 @@ py==1.11.0 # via retry pyarrow==6.0.1 # via flytekit +pycparser==2.21 + # via cffi pyparsing==3.0.7 # via # matplotlib diff --git a/tests/flytekit/unit/core/test_local_cache.py b/tests/flytekit/unit/core/test_local_cache.py index bf4e4c2edb..3f3e56de88 100644 --- a/tests/flytekit/unit/core/test_local_cache.py +++ b/tests/flytekit/unit/core/test_local_cache.py @@ -1,12 +1,15 @@ import datetime import typing from dataclasses import dataclass +from typing import List import pandas from dataclasses_json import dataclass_json from pytest import fixture +from typing_extensions import Annotated -from flytekit import SQLTask, kwtypes +from flytekit import SQLTask, dynamic, kwtypes +from flytekit.core.hash import HashMethod from flytekit.core.local_cache import LocalTaskCache from flytekit.core.task import TaskMetadata, task from flytekit.core.testing import task_mock @@ -252,13 +255,131 @@ def my_wf(a: int, b: str) -> (int, str): assert n_cached_task_calls == 2 -""" -Update SD transformer so that it can to_python_value a Schema literal - - If a Schema literal is detected, copy the uri and use the new decoder to unwrap the uri -Update FS transformer so that it can to_python_value a StructuredDataset literal - - If a StructuredDataset literal is detected, use the uri from that instead. +def test_set_integer_literal_hash_is_not_cached(): + """ + Test to confirm that the local cache is not set in the case of integers, even if we + return an annotated integer. In order to make this very explicit, we define a constant hash + function, i.e. the same value is returned by it regardless of the input. + """ + + def constant_hash_function(a: int) -> str: + return "hash" + + @task + def t0(a: int) -> Annotated[int, HashMethod(function=constant_hash_function)]: + return a + + @task(cache=True, cache_version="0.0.1") + def t1(cached_a: int) -> int: + global n_cached_task_calls + n_cached_task_calls += 1 + return cached_a + + @workflow + def wf(a: int) -> int: + annotated_a = t0(a=a) + return t1(cached_a=annotated_a) + + assert n_cached_task_calls == 0 + assert wf(a=3) == 3 + assert n_cached_task_calls == 1 + # Confirm that the value is not cached, even though we set a hash function that + # returns a constant value and that the task has only one input. + assert wf(a=2) == 2 + assert n_cached_task_calls == 2 + # Confirm that the cache is hit if we execute the workflow with the same value as previous run. + assert wf(a=2) == 2 + assert n_cached_task_calls == 2 + + +def test_pass_annotated_to_downstream_tasks(): + @task + def t0(a: int) -> Annotated[int, HashMethod(function=str)]: + return a + 1 + + @task(cache=True, cache_version="42") + def downstream_t(a: int) -> int: + global n_cached_task_calls + n_cached_task_calls += 1 + return a + 2 + + @dynamic + def t1(a: int) -> int: + v = t0(a=a) + + # We should have a cache miss in the first call to downstream_t and have a cache hit + # on the second call. + v_1 = downstream_t(a=v) + v_2 = downstream_t(a=v) + + return v_1 + v_2 + + assert n_cached_task_calls == 0 + assert t1(a=3) == (6 + 6) + assert n_cached_task_calls == 1 + + +def test_pandas_dataframe_hash(): + """ + Test that cache is hit in the case of pandas dataframes where we annotated dataframes to hash + the contents of the dataframes. + """ + + def hash_pandas_dataframe(df: pandas.DataFrame) -> str: + return str(pandas.util.hash_pandas_object(df)) + + @task + def uncached_data_reading_task() -> Annotated[pandas.DataFrame, HashMethod(hash_pandas_dataframe)]: + return pandas.DataFrame({"column_1": [1, 2, 3]}) + + @task(cache=True, cache_version="0.1") + def cached_data_processing_task(data: pandas.DataFrame) -> pandas.DataFrame: + global n_cached_task_calls + n_cached_task_calls += 1 + return data * 2 + + @workflow + def my_workflow(): + raw_data = uncached_data_reading_task() + cached_data_processing_task(data=raw_data) + + assert n_cached_task_calls == 0 + my_workflow() + assert n_cached_task_calls == 1 + + # Confirm that we see a cache hit in the case of annotated dataframes. + my_workflow() + assert n_cached_task_calls == 1 -Update all plugins that can take in a FlyteSchema to also be able to take in a StructuredDataset. -All tests should work with the presence of SD imports. -""" +def test_list_of_pandas_dataframe_hash(): + """ + Test that cache is hit in the case of a list of pandas dataframes where we annotated dataframes to hash + the contents of the dataframes. + """ + + def hash_pandas_dataframe(df: pandas.DataFrame) -> str: + return str(pandas.util.hash_pandas_object(df)) + + @task + def uncached_data_reading_task() -> List[Annotated[pandas.DataFrame, HashMethod(hash_pandas_dataframe)]]: + return [pandas.DataFrame({"column_1": [1, 2, 3]}), pandas.DataFrame({"column_1": [10, 20, 30]})] + + @task(cache=True, cache_version="0.1") + def cached_data_processing_task(data: List[pandas.DataFrame]) -> List[pandas.DataFrame]: + global n_cached_task_calls + n_cached_task_calls += 1 + return [df * 2 for df in data] + + @workflow + def my_workflow(): + raw_data = uncached_data_reading_task() + cached_data_processing_task(data=raw_data) + + assert n_cached_task_calls == 0 + my_workflow() + assert n_cached_task_calls == 1 + + # Confirm that we see a cache hit in the case of annotated dataframes. + my_workflow() + assert n_cached_task_calls == 1 diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index e511cceb6f..b97d77eef4 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -17,11 +17,15 @@ from marshmallow_enum import LoadDumpOptions from marshmallow_jsonschema import JSONSchema from pandas._testing import assert_frame_equal +from typing_extensions import Annotated import flytekit.common.exceptions.user as user_exceptions from flytekit import kwtypes from flytekit.core.annotation import FlyteAnnotation from flytekit.core.context_manager import FlyteContext, FlyteContextManager +from flytekit.core.dynamic_workflow_task import dynamic +from flytekit.core.hash import HashMethod +from flytekit.core.task import task from flytekit.core.type_engine import ( DataclassTransformer, DictTransformer, @@ -47,6 +51,7 @@ from flytekit.types.pickle import FlytePickle from flytekit.types.pickle.pickle import FlytePickleTransformer from flytekit.types.schema import FlyteSchema +from flytekit.types.schema.types_pandas import PandasDataFrameTransformer from flytekit.types.structured.structured_dataset import StructuredDataset try: @@ -1072,6 +1077,77 @@ def test_dict_to_literal_map_with_wrong_input_type(): TypeEngine.dict_to_literal_map(ctx, input, guessed_python_types) +def test_pass_annotated_to_downstream_tasks(): + """ + Test to confirm that the loaded dataframe is not affected and can be used in @dynamic. + """ + # pandas dataframe hash function + def hash_pandas_dataframe(df: pd.DataFrame) -> str: + return str(pd.util.hash_pandas_object(df)) + + @task + def t0(a: int) -> Annotated[int, HashMethod(function=str)]: + return a + 1 + + @task + def annotated_return_task() -> Annotated[pd.DataFrame, HashMethod(hash_pandas_dataframe)]: + return pd.DataFrame({"column_1": [1, 2, 3]}) + + @task(cache=True, cache_version="42") + def downstream_t(a: int, df: pd.DataFrame) -> int: + return a + 2 + len(df) + + @dynamic + def t1(a: int) -> int: + v = t0(a=a) + df = annotated_return_task() + + # We should have a cache miss in the first call to downstream_t + v_1 = downstream_t(a=v, df=df) + v_2 = downstream_t(a=v, df=df) + + return v_1 + v_2 + + assert t1(a=3) == (6 + 6 + 6) + + +def test_literal_hash_int_not_set(): + """ + Test to confirm that annotating an integer with `HashMethod` does not force the literal to have its + hash set. + """ + ctx = FlyteContext.current_context() + lv = TypeEngine.to_literal( + ctx, 42, Annotated[int, HashMethod(str)], LiteralType(simple=model_types.SimpleType.INTEGER) + ) + assert lv.scalar.primitive.integer == 42 + assert lv.hash is None + + +def test_literal_hash_to_python_value(): + """ + Test to confirm that literals can be converted to python values, regardless of the hash value set in the literal. + """ + ctx = FlyteContext.current_context() + + def constant_hash(df: pd.DataFrame) -> str: + return "h4Sh" + + df = pd.DataFrame(data={"col1": [1, 2], "col2": [3, 4]}) + pandas_df_transformer = PandasDataFrameTransformer() + literal_with_hash_set = TypeEngine.to_literal( + ctx, + df, + Annotated[pd.DataFrame, HashMethod(constant_hash)], + pandas_df_transformer.get_literal_type(pd.DataFrame), + ) + assert literal_with_hash_set.hash == "h4Sh" + # Confirm tha the loaded dataframe is not affected + python_df = TypeEngine.to_python_value(ctx, literal_with_hash_set, pd.DataFrame) + expected_df = pd.DataFrame(data={"col1": [1, 2], "col2": [3, 4]}) + assert expected_df.equals(python_df) + + def test_annotated_simple_types(): def _check_annotation(t, annotation): lt = TypeEngine.to_literal_type(t) diff --git a/tests/flytekit/unit/core/test_type_hints.py b/tests/flytekit/unit/core/test_type_hints.py index 67247f53ee..d256f26c53 100644 --- a/tests/flytekit/unit/core/test_type_hints.py +++ b/tests/flytekit/unit/core/test_type_hints.py @@ -16,6 +16,7 @@ from dataclasses_json import dataclass_json from google.protobuf.struct_pb2 import Struct from pandas._testing import assert_frame_equal +from typing_extensions import Annotated import flytekit from flytekit import ContainerTask, Secret, SQLTask, dynamic, kwtypes, map_task @@ -23,6 +24,7 @@ from flytekit.core.condition import conditional from flytekit.core.context_manager import ExecutionState, FastSerializationSettings, Image, ImageConfig from flytekit.core.data_persistence import FileAccessProvider +from flytekit.core.hash import HashMethod from flytekit.core.node import Node from flytekit.core.promise import NodeOutput, Promise, VoidPromise from flytekit.core.resources import Resources @@ -1780,3 +1782,110 @@ def wf(a: int) -> str: assert wf(a=-10) == "MyInt -10" del TypeEngine._REGISTRY[MyInt] + + +def test_task_annotate_primitive_type_has_no_effect(): + @task + def plus_two( + a: int, + ) -> Annotated[int, HashMethod(str)]: # Note the use of `str` as the hash function for ints. This has no effect. + return a + 2 + + assert plus_two(a=1) == 3 + + ctx = context_manager.FlyteContextManager.current_context() + output_lm = plus_two.dispatch_execute( + ctx, + _literal_models.LiteralMap( + literals={ + "a": _literal_models.Literal( + scalar=_literal_models.Scalar(primitive=_literal_models.Primitive(integer=3)) + ) + } + ), + ) + assert output_lm.literals["o0"].scalar.primitive.integer == 5 + assert output_lm.literals["o0"].hash is None + + +def test_task_hash_return_pandas_dataframe(): + constant_value = "road-hash" + + def constant_function(df: pandas.DataFrame) -> str: + return constant_value + + @task + def t0() -> Annotated[pandas.DataFrame, HashMethod(constant_function)]: + return pandas.DataFrame(data={"col1": [1, 2], "col2": [3, 4]}) + + ctx = context_manager.FlyteContextManager.current_context() + output_lm = t0.dispatch_execute(ctx, _literal_models.LiteralMap(literals={})) + assert output_lm.literals["o0"].hash == constant_value + + # Confirm that the literal containing a hash does not have any effect on the scalar. + df = TypeEngine.to_python_value(ctx, output_lm.literals["o0"], pandas.DataFrame) + expected_df = pandas.DataFrame(data={"col1": [1, 2], "col2": [3, 4]}) + assert df.equals(expected_df) + + +def test_workflow_containing_multiple_annotated_tasks(): + def hash_function_t0(df: pandas.DataFrame) -> str: + return "hash-0" + + @task + def t0() -> Annotated[pandas.DataFrame, HashMethod(hash_function_t0)]: + return pandas.DataFrame(data={"col1": [1, 2], "col2": [3, 4]}) + + def hash_function_t1(df: pandas.DataFrame) -> str: + return "hash-1" + + @task + def t1() -> Annotated[pandas.DataFrame, HashMethod(hash_function_t1)]: + return pandas.DataFrame(data={"col1": [10, 20], "col2": [30, 40]}) + + @task + def t2() -> pandas.DataFrame: + return pandas.DataFrame(data={"col1": [100, 200], "col2": [300, 400]}) + + # Auxiliary task used to sum up the dataframes. It demonstrates that the use of `Annotated` does not + # have any impact in the definition and execution of cached or uncached downstream tasks + @task + def sum_dataframes(df0: pandas.DataFrame, df1: pandas.DataFrame, df2: pandas.DataFrame) -> pandas.DataFrame: + return df0 + df1 + df2 + + @workflow + def wf() -> pandas.DataFrame: + df0 = t0() + df1 = t1() + df2 = t2() + return sum_dataframes(df0=df0, df1=df1, df2=df2) + + df = wf() + + expected_df = pandas.DataFrame(data={"col1": [1 + 10 + 100, 2 + 20 + 200], "col2": [3 + 30 + 300, 4 + 40 + 400]}) + assert expected_df.equals(df) + + +def test_list_containing_multiple_annotated_pandas_dataframes(): + def hash_pandas_dataframe(df: pandas.DataFrame) -> str: + return str(pandas.util.hash_pandas_object(df)) + + @task + def produce_list_of_annotated_dataframes() -> typing.List[ + Annotated[pandas.DataFrame, HashMethod(hash_pandas_dataframe)] + ]: + return [pandas.DataFrame({"column_1": [1, 2, 3]}), pandas.DataFrame({"column_1": [4, 5, 6]})] + + @task(cache=True, cache_version="v0") + def sum_list_of_pandas_dataframes(lst: typing.List[pandas.DataFrame]) -> pandas.DataFrame: + return sum(lst) + + @workflow + def wf() -> pandas.DataFrame: + lst = produce_list_of_annotated_dataframes() + return sum_list_of_pandas_dataframes(lst=lst) + + df = wf() + + expected_df = pandas.DataFrame({"column_1": [5, 7, 9]}) + assert expected_df.equals(df) diff --git a/tests/flytekit_compatibility/test_structured_dataset.py b/tests/flytekit_compatibility/test_structured_dataset.py index 20965cb802..f935305210 100644 --- a/tests/flytekit_compatibility/test_structured_dataset.py +++ b/tests/flytekit_compatibility/test_structured_dataset.py @@ -1,8 +1,12 @@ import pandas as pd +from flytekit.configuration.sdk import USE_STRUCTURED_DATASET from flytekit.core.type_engine import TypeEngine def test_pandas_is_schema_with_flag(): + # This test can only be run iff USE_STRUCTURED_DATASET is not set + assert USE_STRUCTURED_DATASET.get() is False + lt = TypeEngine.to_literal_type(pd.DataFrame) assert lt.schema is not None From e9e2b49e1ac4f79e874e117672ae70108c31e786 Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Wed, 2 Mar 2022 16:13:51 -0800 Subject: [PATCH 103/128] FlyteRemote fetch of conditional nodes (#772) Signed-off-by: Yee Hing Tong Signed-off-by: maximsmol --- flytekit/clis/flyte_cli/main.py | 1 - flytekit/models/core/workflow.py | 2 +- flytekit/remote/component_nodes.py | 35 +++++++++++++++++++++++++++++- flytekit/remote/launch_plan.py | 7 ++++++ flytekit/remote/nodes.py | 28 ++++++++++++++++++------ flytekit/remote/remote.py | 34 ++++++++++++++++++----------- 6 files changed, 84 insertions(+), 23 deletions(-) diff --git a/flytekit/clis/flyte_cli/main.py b/flytekit/clis/flyte_cli/main.py index 82300a9c7f..6415e92b15 100644 --- a/flytekit/clis/flyte_cli/main.py +++ b/flytekit/clis/flyte_cli/main.py @@ -314,7 +314,6 @@ def _render_schedule_expr(lp): ) _insecure_option = _click.option(*_INSECURE_FLAGS, is_flag=True, help="Do not use SSL") _urn_option = _click.option("-u", "--urn", required=True, help="The unique identifier for an entity.") - _optional_urn_option = _click.option("-u", "--urn", required=False, help="The unique identifier for an entity.") _host_option = _click.option( diff --git a/flytekit/models/core/workflow.py b/flytekit/models/core/workflow.py index 6338a6695f..e9e68ce485 100644 --- a/flytekit/models/core/workflow.py +++ b/flytekit/models/core/workflow.py @@ -99,7 +99,7 @@ def else_node(self): def error(self): """ An error to throw in case none of the branches were taken. - :rtype: flytekit.models.core.errors.ContainerError + :rtype: flytekit.models.types.Error """ return self._error diff --git a/flytekit/remote/component_nodes.py b/flytekit/remote/component_nodes.py index 877a6d6494..2d6acca1ff 100644 --- a/flytekit/remote/component_nodes.py +++ b/flytekit/remote/component_nodes.py @@ -41,7 +41,7 @@ def promote_from_model( if base_model.reference_id in tasks: task = tasks[base_model.reference_id] - _logging.info(f"Found existing task template for {task.id}, will not retrieve from Admin") + _logging.debug(f"Found existing task template for {task.id}, will not retrieve from Admin") flyte_task = FlyteTask.promote_from_model(task) return cls(flyte_task) @@ -124,3 +124,36 @@ def promote_from_model( raise _system_exceptions.FlyteSystemException( "Bad workflow node model, neither subworkflow nor launchplan specified." ) + + +class FlyteBranchNode(_workflow_model.BranchNode): + def __init__(self, if_else: _workflow_model.IfElseBlock): + super().__init__(if_else) + + @classmethod + def promote_from_model( + cls, + base_model: _workflow_model.BranchNode, + sub_workflows: Dict[id_models.Identifier, _workflow_model.WorkflowTemplate], + node_launch_plans: Dict[id_models.Identifier, _launch_plan_model.LaunchPlanSpec], + tasks: Dict[id_models.Identifier, _task_model.TaskTemplate], + ) -> "FlyteBranchNode": + + from flytekit.remote.nodes import FlyteNode + + block = base_model.if_else + + else_node = None + if block.else_node: + else_node = FlyteNode.promote_from_model(block.else_node, sub_workflows, node_launch_plans, tasks) + + block.case._then_node = FlyteNode.promote_from_model( + block.case.then_node, sub_workflows, node_launch_plans, tasks + ) + + for o in block.other: + o._then_node = FlyteNode.promote_from_model(o.then_node, sub_workflows, node_launch_plans, tasks) + + new_if_else_block = _workflow_model.IfElseBlock(block.case, block.other, else_node, block.error) + + return cls(new_if_else_block) diff --git a/flytekit/remote/launch_plan.py b/flytekit/remote/launch_plan.py index 32b3a93aec..4ad7272957 100644 --- a/flytekit/remote/launch_plan.py +++ b/flytekit/remote/launch_plan.py @@ -28,6 +28,13 @@ def __init__(self, id, *args, **kwargs): def name(self) -> str: return self._name + # If fetched when creating this object, can store it here. + self._flyte_workflow = None + + @property + def flyte_workflow(self) -> Optional["FlyteWorkflow"]: + return self._flyte_workflow + @classmethod def promote_from_model( cls, id: id_models.Identifier, model: _launch_plan_models.LaunchPlanSpec diff --git a/flytekit/remote/nodes.py b/flytekit/remote/nodes.py index 4131905c5b..e7ee0dc487 100644 --- a/flytekit/remote/nodes.py +++ b/flytekit/remote/nodes.py @@ -27,15 +27,19 @@ def __init__( flyte_task: Optional["FlyteTask"] = None, flyte_workflow: Optional["FlyteWorkflow"] = None, flyte_launch_plan: Optional["FlyteLaunchPlan"] = None, - flyte_branch=None, + flyte_branch_node: Optional["FlyteBranchNode"] = None, ): - non_none_entities = list(filter(None, [flyte_task, flyte_workflow, flyte_launch_plan, flyte_branch])) + # todo: flyte_branch_node is the only non-entity here, feels wrong, it should probably be a Condition + # or the other ones changed. + non_none_entities = list(filter(None, [flyte_task, flyte_workflow, flyte_launch_plan, flyte_branch_node])) if len(non_none_entities) != 1: raise _user_exceptions.FlyteAssertion( "An Flyte node must have one underlying entity specified at once. Received the following " "entities: {}".format(non_none_entities) ) - self._flyte_entity = flyte_task or flyte_workflow or flyte_launch_plan or flyte_branch + # todo: wip - flyte_branch_node is a hack, it should be a Condition, but backing out a Condition object from + # the compiled IfElseBlock is cumbersome, shouldn't do it if we can get away with it. + self._flyte_entity = flyte_task or flyte_workflow or flyte_launch_plan or flyte_branch_node workflow_node = None if flyte_workflow is not None: @@ -46,7 +50,6 @@ def __init__( task_node = None if flyte_task: task_node = _component_nodes.FlyteTaskNode(flyte_task) - branch_node = None super(FlyteNode, self).__init__( id=id, @@ -56,7 +59,7 @@ def __init__( output_aliases=[], task_node=task_node, workflow_node=workflow_node, - branch_node=branch_node, + branch_node=flyte_branch_node, ) self._upstream = upstream_nodes @@ -78,7 +81,7 @@ def promote_from_model( _logging.warning(f"Should not call promote from model on a start node or end node {model}") return None - flyte_task_node, flyte_workflow_node = None, None + flyte_task_node, flyte_workflow_node, flyte_branch_node = None, None, None if model.task_node is not None: flyte_task_node = _component_nodes.FlyteTaskNode.promote_from_model(model.task_node, tasks) elif model.workflow_node is not None: @@ -88,7 +91,10 @@ def promote_from_model( node_launch_plans, tasks, ) - # TODO: Implement branch node https://github.com/flyteorg/flyte/issues/1116 + elif model.branch_node is not None: + flyte_branch_node = _component_nodes.FlyteBranchNode.promote_from_model( + model.branch_node, sub_workflows, node_launch_plans, tasks + ) else: raise _system_exceptions.FlyteSystemException( f"Bad Node model, neither task nor workflow detected, node: {model}" @@ -132,6 +138,14 @@ def promote_from_model( raise _system_exceptions.FlyteSystemException( "Bad FlyteWorkflowNode model, both launch plan and workflow are None" ) + elif flyte_branch_node is not None: + return cls( + id=node_model_id, + upstream_nodes=[], # set downstream, model doesn't contain this information + bindings=model.inputs, + metadata=model.metadata, + flyte_branch_node=flyte_branch_node, + ) raise _system_exceptions.FlyteSystemException("Bad FlyteNode model, both task and workflow nodes are empty") @property diff --git a/flytekit/remote/remote.py b/flytekit/remote/remote.py index 93adc04db8..50d29e8e23 100644 --- a/flytekit/remote/remote.py +++ b/flytekit/remote/remote.py @@ -440,6 +440,7 @@ def fetch_launch_plan( wf_id = flyte_launch_plan.workflow_id workflow = self.fetch_workflow(wf_id.project, wf_id.domain, wf_id.name, wf_id.version) flyte_launch_plan._interface = workflow.interface + flyte_launch_plan._flyte_workflow = workflow flyte_launch_plan.guessed_python_interface = Interface( inputs=TypeEngine.guess_python_types(flyte_launch_plan.interface.inputs), outputs=TypeEngine.guess_python_types(flyte_launch_plan.interface.outputs), @@ -1053,6 +1054,7 @@ def sync_workflow_execution( if execution.spec.launch_plan.resource_type == ResourceType.TASK: # This condition is only true for single-task executions flyte_entity = self.fetch_task(lp_id.project, lp_id.domain, lp_id.name, lp_id.version) + node_interface = flyte_entity.interface if sync_nodes: # Need to construct the mapping. There should've been returned exactly three nodes, a start, # an end, and a task node. @@ -1080,10 +1082,10 @@ def sync_workflow_execution( ) else: # This is the default case, an execution of a normal workflow through a launch plan - wf_id = self.fetch_launch_plan(lp_id.project, lp_id.domain, lp_id.name, lp_id.version).workflow_id - flyte_entity = self.fetch_workflow(wf_id.project, wf_id.domain, wf_id.name, wf_id.version) - execution._flyte_workflow = flyte_entity - node_mapping = flyte_entity._node_map + fetched_lp = self.fetch_launch_plan(lp_id.project, lp_id.domain, lp_id.name, lp_id.version) + node_interface = fetched_lp.flyte_workflow.interface + execution._flyte_workflow = fetched_lp.flyte_workflow + node_mapping = fetched_lp.flyte_workflow._node_map # update node executions (if requested), and inputs/outputs if sync_nodes: @@ -1091,10 +1093,12 @@ def sync_workflow_execution( for n in underlying_node_executions: node_execs[n.id.node_id] = self.sync_node_execution(n, node_mapping) execution._node_executions = node_execs - return self._assign_inputs_and_outputs(execution, execution_data, flyte_entity.interface) + return self._assign_inputs_and_outputs(execution, execution_data, node_interface) def sync_node_execution( - self, execution: FlyteNodeExecution, node_mapping: typing.Dict[str, FlyteNode] + self, + execution: FlyteNodeExecution, + node_mapping: typing.Dict[str, FlyteNode], ) -> FlyteNodeExecution: """ Get data backing a node execution. These FlyteNodeExecution objects should've come from Admin with the model @@ -1201,13 +1205,9 @@ def sync_node_execution( for t in iterate_task_executions(self.client, execution.id) ] execution._interface = dynamic_flyte_wf.interface - else: - # If it does not, then it should be a static subworkflow - if not isinstance(execution._node.flyte_entity, FlyteWorkflow): - remote_logger.error( - f"NE {execution} entity should be a workflow, {type(execution._node)}, {execution._node}" - ) - raise Exception(f"Node entity has type {type(execution._node)}") + + # Handle the case where it's a static subworkflow + elif isinstance(execution._node.flyte_entity, FlyteWorkflow): sub_flyte_workflow = execution._node.flyte_entity sub_node_mapping = {n.id: n for n in sub_flyte_workflow.flyte_nodes} execution._underlying_node_executions = [ @@ -1216,6 +1216,14 @@ def sync_node_execution( ] execution._interface = sub_flyte_workflow.interface + # Handle the case where it's a branch node + elif execution._node.branch_node is not None: + remote_logger.debug("Skipping remote node execution for now") + return execution + else: + remote_logger.error(f"NE {execution} undeterminable, {type(execution._node)}, {execution._node}") + raise Exception(f"Node execution undeterminable, entity has type {type(execution._node)}") + # This is the plain ol' task execution case else: execution._task_executions = [ From ef0291d888ca9a9e9ce07674b5b074ad951057a4 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Fri, 4 Mar 2022 01:53:01 +0800 Subject: [PATCH 104/128] Removed root logger (#871) * Removed root logger Signed-off-by: Kevin Su * Updated logger Signed-off-by: Kevin Su * Fixed lint Signed-off-by: Kevin Su * Updated log level Signed-off-by: Kevin Su * Updated logger Signed-off-by: Kevin Su * Updated logger Signed-off-by: Kevin Su * Updated logger Signed-off-by: Kevin Su * lint fixed Signed-off-by: Kevin Su Signed-off-by: maximsmol --- flytekit/bin/entrypoint.py | 4 ++-- flytekit/clients/raw.py | 3 +-- flytekit/clis/sdk_in_container/pyflyte.py | 4 ++-- flytekit/clis/sdk_in_container/serialize.py | 7 ++++--- flytekit/configuration/__init__.py | 7 ++++--- flytekit/core/base_task.py | 2 +- flytekit/core/context_manager.py | 12 ++++++------ flytekit/core/interface.py | 5 ++--- flytekit/core/mock_stats.py | 5 +++-- flytekit/core/reference_entity.py | 2 +- flytekit/core/shim_task.py | 2 +- flytekit/core/tracker.py | 10 +++++----- flytekit/core/utils.py | 6 +++--- flytekit/extras/tasks/shell.py | 6 +++--- flytekit/loggers.py | 2 ++ flytekit/remote/component_nodes.py | 4 ++-- flytekit/remote/nodes.py | 4 ++-- flytekit/remote/remote.py | 9 ++++----- flytekit/tools/subprocess.py | 7 ++++--- .../flytekitplugins/awssagemaker/training.py | 8 ++++---- plugins/flytekit-dolt/flytekitplugins/dolt/schema.py | 3 --- .../flytekitplugins/great_expectations/schema.py | 4 ++-- .../flytekitplugins/great_expectations/task.py | 4 ++-- .../flytekitplugins/papermill/task.py | 4 ++-- 24 files changed, 62 insertions(+), 62 deletions(-) diff --git a/flytekit/bin/entrypoint.py b/flytekit/bin/entrypoint.py index 9d07d38a63..5123e90225 100644 --- a/flytekit/bin/entrypoint.py +++ b/flytekit/bin/entrypoint.py @@ -1,6 +1,5 @@ import contextlib import datetime as _datetime -import logging as python_logging import os as _os import pathlib import traceback as _traceback @@ -31,6 +30,7 @@ from flytekit.exceptions import scopes as _scoped_exceptions from flytekit.exceptions import scopes as _scopes from flytekit.interfaces.stats.taggable import get_stats as _get_stats +from flytekit.loggers import entrypoint_logger from flytekit.loggers import entrypoint_logger as logger from flytekit.models import dynamic_job as _dynamic_job from flytekit.models import literals as _literal_models @@ -209,7 +209,7 @@ def setup_execution( "api_version": _api_version, }, ), - logging=python_logging, + logging=entrypoint_logger, tmp_dir=user_workspace_dir, raw_output_prefix=ctx.file_access._raw_output_prefix, checkpoint=checkpointer, diff --git a/flytekit/clients/raw.py b/flytekit/clients/raw.py index 24d3e118b3..9754375dec 100644 --- a/flytekit/clients/raw.py +++ b/flytekit/clients/raw.py @@ -1,7 +1,6 @@ from __future__ import annotations import base64 as _base64 -import logging as _logging import subprocess import time from typing import Optional @@ -823,7 +822,7 @@ def get_token(token_endpoint, authorization_header, scope): body["scope"] = scope response = _requests.post(token_endpoint, data=body, headers=headers) if response.status_code != 200: - _logging.error("Non-200 ({}) received from IDP: {}".format(response.status_code, response.text)) + cli_logger.error("Non-200 ({}) received from IDP: {}".format(response.status_code, response.text)) raise FlyteAuthenticationException("Non-200 received from IDP") response = response.json() diff --git a/flytekit/clis/sdk_in_container/pyflyte.py b/flytekit/clis/sdk_in_container/pyflyte.py index e0d3859122..56b8328701 100644 --- a/flytekit/clis/sdk_in_container/pyflyte.py +++ b/flytekit/clis/sdk_in_container/pyflyte.py @@ -1,4 +1,3 @@ -import logging as _logging import os as _os from pathlib import Path @@ -16,6 +15,7 @@ from flytekit.configuration.internal import CONFIGURATION_PATH from flytekit.configuration.platform import URL as _URL from flytekit.configuration.sdk import WORKFLOW_PACKAGES as _WORKFLOW_PACKAGES +from flytekit.loggers import cli_logger def validate_package(ctx, param, values): @@ -61,7 +61,7 @@ def main(ctx, config=None, pkgs=None, insecure=None): # Update the logger if it's set log_level = _internal_config.LOGGING_LEVEL.get() or _sdk_config.LOGGING_LEVEL.get() if log_level is not None: - _logging.getLogger().setLevel(log_level) + cli_logger.getLogger().setLevel(log_level) ctx.obj = dict() diff --git a/flytekit/clis/sdk_in_container/serialize.py b/flytekit/clis/sdk_in_container/serialize.py index 5fd68aebc4..4668b34b1d 100644 --- a/flytekit/clis/sdk_in_container/serialize.py +++ b/flytekit/clis/sdk_in_container/serialize.py @@ -1,4 +1,4 @@ -import logging as _logging +import logging import math as _math import os as _os import sys @@ -21,6 +21,7 @@ from flytekit.core.workflow import WorkflowBase from flytekit.exceptions.scopes import system_entry_point from flytekit.exceptions.user import FlyteValidationException +from flytekit.loggers import cli_logger from flytekit.models import launch_plan as _launch_plan_models from flytekit.models import task as task_models from flytekit.models.admin import workflow as admin_workflow_models @@ -294,7 +295,7 @@ def serialize(ctx, image, local_source_root, in_container_config_path, in_contai @click.option("-f", "--folder", type=click.Path(exists=True)) @click.pass_context def workflows(ctx, folder=None): - _logging.getLogger().setLevel(_logging.DEBUG) + cli_logger.getLogger().setLevel(logging.DEBUG) if folder: click.echo(f"Writing output to {folder}") @@ -322,7 +323,7 @@ def fast(ctx): @click.option("-f", "--folder", type=click.Path(exists=True)) @click.pass_context def fast_workflows(ctx, folder=None): - _logging.getLogger().setLevel(_logging.DEBUG) + cli_logger.getLogger().setLevel(logging.DEBUG) if folder: click.echo(f"Writing output to {folder}") diff --git a/flytekit/configuration/__init__.py b/flytekit/configuration/__init__.py index 6dd9c28d35..3f36992aad 100644 --- a/flytekit/configuration/__init__.py +++ b/flytekit/configuration/__init__.py @@ -1,7 +1,8 @@ -import logging as _logging import os as _os import pathlib as _pathlib +from flytekit.loggers import logger + def set_flyte_config_file(config_file_path): """ @@ -14,12 +15,12 @@ def set_flyte_config_file(config_file_path): original_config_file_path = config_file_path config_file_path = _os.path.abspath(config_file_path) if not _pathlib.Path(config_file_path).is_file(): - _logging.warning( + logger.warning( f"No config file provided or invalid flyte config_file_path {original_config_file_path} specified." ) _os.environ[_internal.CONFIGURATION_PATH.env_var] = config_file_path elif _internal.CONFIGURATION_PATH.env_var in _os.environ: - _logging.debug("Deleting configuration path {} from env".format(_internal.CONFIGURATION_PATH.env_var)) + logger.debug("Deleting configuration path {} from env".format(_internal.CONFIGURATION_PATH.env_var)) del _os.environ[_internal.CONFIGURATION_PATH.env_var] _common.CONFIGURATION_SINGLETON.reset_config(config_file_path) diff --git a/flytekit/core/base_task.py b/flytekit/core/base_task.py index cb540420db..0ca49a9169 100644 --- a/flytekit/core/base_task.py +++ b/flytekit/core/base_task.py @@ -480,7 +480,7 @@ def dispatch_execute( logger.exception(f"Exception when executing {e}") raise e - logger.info(f"Task executed successfully in user level, outputs: {native_outputs}") + logger.debug("Task executed successfully in user level") # Lets run the post_execute method. This may result in a IgnoreOutputs Exception, which is # bubbled up to be handled at the callee layer. native_outputs = self.post_execute(new_user_params, native_outputs) diff --git a/flytekit/core/context_manager.py b/flytekit/core/context_manager.py index b57ff043b5..e1f77cb123 100644 --- a/flytekit/core/context_manager.py +++ b/flytekit/core/context_manager.py @@ -14,7 +14,6 @@ from __future__ import annotations import datetime as _datetime -import logging import logging as _logging import os import pathlib @@ -40,6 +39,7 @@ from flytekit.core.node import Node from flytekit.interfaces.cli_identifiers import WorkflowExecutionIdentifier from flytekit.interfaces.stats import taggable +from flytekit.loggers import logger, user_space_logger from flytekit.models.core import identifier as _identifier # TODO: resolve circular import from flytekit.core.python_auto_container import TaskResolverMixin @@ -157,7 +157,7 @@ class ExecutionParameters(object): class Builder(object): stats: taggable.TaggableStats execution_date: datetime - logging: _logging + logging: _logging.Logger execution_id: str attrs: typing.Dict[str, typing.Any] working_dir: typing.Union[os.PathLike, utils.AutoDeletingTempDir] @@ -246,7 +246,7 @@ def stats(self) -> taggable.TaggableStats: return self._stats @property - def logging(self) -> _logging: + def logging(self) -> _logging.Logger: """ A handle to a useful logging object. TODO: Usage examples @@ -899,7 +899,7 @@ def push_context(ctx: FlyteContext, f: Optional[traceback.FrameSummary] = None) ctx.set_stackframe(f) FlyteContextManager._OBJS.append(ctx) t = "\t" - logging.debug( + logger.debug( f"{t * ctx.level}[{len(FlyteContextManager._OBJS)}] Pushing context - {'compile' if ctx.compilation_state else 'execute'}, branch[{ctx.in_a_condition}], {ctx.get_origin_stackframe_repr()}" ) return ctx @@ -908,7 +908,7 @@ def push_context(ctx: FlyteContext, f: Optional[traceback.FrameSummary] = None) def pop_context() -> FlyteContext: ctx = FlyteContextManager._OBJS.pop() t = "\t" - logging.debug( + logger.debug( f"{t * ctx.level}[{len(FlyteContextManager._OBJS) + 1}] Popping context - {'compile' if ctx.compilation_state else 'execute'}, branch[{ctx.in_a_condition}], {ctx.get_origin_stackframe_repr()}" ) if len(FlyteContextManager._OBJS) == 0: @@ -963,7 +963,7 @@ def initialize(): execution_id=str(WorkflowExecutionIdentifier.promote_from_model(default_execution_id)), execution_date=_datetime.datetime.utcnow(), stats=mock_stats.MockStats(), - logging=_logging, + logging=user_space_logger, tmp_dir=user_space_path, raw_output_prefix=default_context.file_access._raw_output_prefix, ) diff --git a/flytekit/core/interface.py b/flytekit/core/interface.py index b158a9434c..fed0793e4a 100644 --- a/flytekit/core/interface.py +++ b/flytekit/core/interface.py @@ -3,7 +3,6 @@ import collections import copy import inspect -import logging as _logging import typing from collections import OrderedDict from typing import Any, Dict, Generator, List, Optional, Tuple, Type, TypeVar, Union @@ -258,7 +257,7 @@ def _change_unrecognized_type_to_pickle(t: Type[T]) -> Type[T]: else: TypeEngine.get_transformer(t) except ValueError: - _logging.warning( + logger.warning( f"Unsupported Type {t} found, Flyte will default to use PickleFile as the transport. " f"Pickle can only be used to send objects between the exact same version of Python, " f"and we strongly recommend to use python type that flyte support." @@ -267,7 +266,7 @@ def _change_unrecognized_type_to_pickle(t: Type[T]) -> Type[T]: return t -def transform_function_to_interface(fn: Callable, docstring: Optional[Docstring] = None) -> Interface: +def transform_function_to_interface(fn: typing.Callable, docstring: Optional[Docstring] = None) -> Interface: """ From the annotations on a task function that the user should have provided, and the output names they want to use for each output parameter, construct the TypedInterface object diff --git a/flytekit/core/mock_stats.py b/flytekit/core/mock_stats.py index e6fe78fe8f..dedce6ae7c 100644 --- a/flytekit/core/mock_stats.py +++ b/flytekit/core/mock_stats.py @@ -1,5 +1,6 @@ import datetime as _datetime -import logging + +from flytekit.loggers import logger class MockStats(object): @@ -25,7 +26,7 @@ def decr(self, metric, count=1, tags=None, **kwargs): self._records_tags[full_name] = tags or {} def timing(self, metric): - logging.warning("mock timing isn't implemented yet.") + logger.warning("mock timing isn't implemented yet.") def timer(self, metric, tags=None, **kwargs): return _Timer(self, metric, tags=tags or {}) diff --git a/flytekit/core/reference_entity.py b/flytekit/core/reference_entity.py index 42be1313cc..77b96e6892 100644 --- a/flytekit/core/reference_entity.py +++ b/flytekit/core/reference_entity.py @@ -119,7 +119,7 @@ def unwrap_literal_map_and_execute( except Exception as e: logger.exception(f"Exception when executing {e}") raise e - logger.info(f"Task executed successfully in user level, outputs: {native_outputs}") + logger.debug("Task executed successfully in user level") expected_output_names = list(self.python_interface.outputs.keys()) if len(expected_output_names) == 1: diff --git a/flytekit/core/shim_task.py b/flytekit/core/shim_task.py index ad6d39ef38..76e396c518 100644 --- a/flytekit/core/shim_task.py +++ b/flytekit/core/shim_task.py @@ -98,7 +98,7 @@ def dispatch_execute( logger.exception(f"Exception when executing {e}") raise e - logger.info(f"Task executed successfully in user level, outputs: {native_outputs}") + logger.debug("Task executed successfully in user level") # Lets run the post_execute method. This may result in a IgnoreOutputs Exception, which is # bubbled up to be handled at the callee layer. native_outputs = self.post_execute(new_user_params, native_outputs) diff --git a/flytekit/core/tracker.py b/flytekit/core/tracker.py index 56f145b4b6..d72adb6b6a 100644 --- a/flytekit/core/tracker.py +++ b/flytekit/core/tracker.py @@ -1,10 +1,10 @@ import importlib as _importlib import inspect import inspect as _inspect -import logging as _logging from typing import Callable from flytekit.exceptions import system as _system_exceptions +from flytekit.loggers import logger class InstanceTrackingMeta(type): @@ -70,12 +70,12 @@ def find_lhs(self) -> str: if self._instantiated_in is None or self._instantiated_in == "": raise _system_exceptions.FlyteSystemException(f"Object {self} does not have an _instantiated in") - _logging.debug(f"Looking for LHS for {self} from {self._instantiated_in}") + logger.debug(f"Looking for LHS for {self} from {self._instantiated_in}") m = _importlib.import_module(self._instantiated_in) for k in dir(m): try: if getattr(m, k) is self: - _logging.debug(f"Found LHS for {self}, {k}") + logger.debug(f"Found LHS for {self}, {k}") self._lhs = k return k except ValueError as err: @@ -84,10 +84,10 @@ def find_lhs(self) -> str: # a.any() or a.all() # Since dataframes aren't registrable entities to begin with we swallow any errors they raise and # continue looping through m. - _logging.warning("Caught ValueError {} while attempting to auto-assign name".format(err)) + logger.warning("Caught ValueError {} while attempting to auto-assign name".format(err)) pass - _logging.error(f"Could not find LHS for {self} in {self._instantiated_in}") + logger.error(f"Could not find LHS for {self} in {self._instantiated_in}") raise _system_exceptions.FlyteSystemException(f"Error looking for LHS in {self._instantiated_in}") diff --git a/flytekit/core/utils.py b/flytekit/core/utils.py index db3d4b573b..7f87c4de91 100644 --- a/flytekit/core/utils.py +++ b/flytekit/core/utils.py @@ -1,4 +1,3 @@ -import logging as _logging import os as _os import shutil as _shutil import tempfile as _tempfile @@ -8,6 +7,7 @@ from typing import Dict, List, Optional from flytekit.configuration import resources as _resource_config +from flytekit.loggers import logger from flytekit.models import task as _task_models @@ -220,14 +220,14 @@ def __init__(self, context_statement): self._start_process_time = None def __enter__(self): - _logging.info("Entering timed context: {}".format(self._context_statement)) + logger.info("Entering timed context: {}".format(self._context_statement)) self._start_wall_time = _time.perf_counter() self._start_process_time = _time.process_time() def __exit__(self, exc_type, exc_val, exc_tb): end_wall_time = _time.perf_counter() end_process_time = _time.process_time() - _logging.info( + logger.info( "Exiting timed context: {} [Wall Time: {}s, Process Time: {}s]".format( self._context_statement, end_wall_time - self._start_wall_time, diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py index 2686c899d8..60f9f691b3 100644 --- a/flytekit/extras/tasks/shell.py +++ b/flytekit/extras/tasks/shell.py @@ -1,5 +1,4 @@ import datetime -import logging import os import string import subprocess @@ -11,6 +10,7 @@ from flytekit.core.interface import Interface from flytekit.core.python_function_task import PythonInstanceTask from flytekit.core.task import TaskPlugins +from flytekit.loggers import logger from flytekit.types.directory import FlyteDirectory from flytekit.types.file import FlyteFile @@ -191,7 +191,7 @@ def execute(self, **kwargs) -> typing.Any: """ Executes the given script by substituting the inputs and outputs and extracts the outputs from the filesystem """ - logging.info(f"Running shell script as type {self.task_type}") + logger.info(f"Running shell script as type {self.task_type}") if self.script_file: with open(self.script_file) as f: self._script = f.read() @@ -212,7 +212,7 @@ def execute(self, **kwargs) -> typing.Any: except subprocess.CalledProcessError as e: files = os.listdir("./") fstr = "\n-".join(files) - logging.error( + logger.error( f"Failed to Execute Script, return-code {e.returncode} \n" f"StdErr: {e.stderr}\n" f"StdOut: {e.stdout}\n" diff --git a/flytekit/loggers.py b/flytekit/loggers.py index bc3e243883..07ac1abdee 100644 --- a/flytekit/loggers.py +++ b/flytekit/loggers.py @@ -30,11 +30,13 @@ "cli": logger.getChild("cli"), "remote": logger.getChild("remote"), "entrypoint": logger.getChild("entrypoint"), + "user_space": logger.getChild("user_space"), } auth_logger = child_loggers["auth"] cli_logger = child_loggers["cli"] remote_logger = child_loggers["remote"] entrypoint_logger = child_loggers["entrypoint"] +user_space_logger = child_loggers["user_space"] # create console handler ch = logging.StreamHandler() diff --git a/flytekit/remote/component_nodes.py b/flytekit/remote/component_nodes.py index 2d6acca1ff..a52738fad7 100644 --- a/flytekit/remote/component_nodes.py +++ b/flytekit/remote/component_nodes.py @@ -1,7 +1,7 @@ -import logging as _logging from typing import Dict from flytekit.exceptions import system as _system_exceptions +from flytekit.loggers import remote_logger from flytekit.models import launch_plan as _launch_plan_model from flytekit.models import task as _task_model from flytekit.models.core import identifier as id_models @@ -41,7 +41,7 @@ def promote_from_model( if base_model.reference_id in tasks: task = tasks[base_model.reference_id] - _logging.debug(f"Found existing task template for {task.id}, will not retrieve from Admin") + remote_logger.debug(f"Found existing task template for {task.id}, will not retrieve from Admin") flyte_task = FlyteTask.promote_from_model(task) return cls(flyte_task) diff --git a/flytekit/remote/nodes.py b/flytekit/remote/nodes.py index e7ee0dc487..0d73678b7e 100644 --- a/flytekit/remote/nodes.py +++ b/flytekit/remote/nodes.py @@ -1,6 +1,5 @@ from __future__ import annotations -import logging as _logging from typing import Dict, List, Optional, Union from flytekit.core import constants as _constants @@ -8,6 +7,7 @@ from flytekit.core.promise import NodeOutput from flytekit.exceptions import system as _system_exceptions from flytekit.exceptions import user as _user_exceptions +from flytekit.loggers import remote_logger from flytekit.models import launch_plan as _launch_plan_model from flytekit.models import task as _task_model from flytekit.models.core import identifier as id_models @@ -78,7 +78,7 @@ def promote_from_model( node_model_id = model.id # TODO: Consider removing if id in {_constants.START_NODE_ID, _constants.END_NODE_ID}: - _logging.warning(f"Should not call promote from model on a start node or end node {model}") + remote_logger.warning(f"Should not call promote from model on a start node or end node {model}") return None flyte_task_node, flyte_workflow_node, flyte_branch_node = None, None, None diff --git a/flytekit/remote/remote.py b/flytekit/remote/remote.py index 50d29e8e23..56b93f6715 100644 --- a/flytekit/remote/remote.py +++ b/flytekit/remote/remote.py @@ -5,7 +5,6 @@ """ from __future__ import annotations -import logging import os import time import typing @@ -636,9 +635,9 @@ def _register_entity_if_not_exists(self, entity: WorkflowBase, resolved_identifi else: raise NotImplementedError(f"We don't support registering this kind of entity: {node.flyte_entity}") except FlyteEntityAlreadyExistsException: - logging.info(f"{entity.name} already exists") + remote_logger.info(f"{entity.name} already exists") except Exception as e: - logging.info(f"Failed to register entity {entity.name} with error {e}") + remote_logger.info(f"Failed to register entity {entity.name} with error {e}") #################### # Execute Entities # @@ -924,7 +923,7 @@ def _( try: flyte_workflow: FlyteWorkflow = self.fetch_workflow(**resolved_identifiers_dict) except FlyteEntityNotExistException: - logging.info("Try to register FlyteWorkflow because it wasn't found in Flyte Admin!") + remote_logger.info("Try to register FlyteWorkflow because it wasn't found in Flyte Admin!") self._register_entity_if_not_exists(entity, resolved_identifiers_dict) flyte_workflow: FlyteWorkflow = self.register(entity, **resolved_identifiers_dict) flyte_workflow.guessed_python_interface = entity.python_interface @@ -933,7 +932,7 @@ def _( try: self.fetch_launch_plan(**resolved_identifiers_dict) except FlyteEntityNotExistException: - logging.info("Try to register default launch plan because it wasn't found in Flyte Admin!") + remote_logger.info("Try to register default launch plan because it wasn't found in Flyte Admin!") default_lp = LaunchPlan.get_default_launch_plan(ctx, entity) self.register(default_lp, **resolved_identifiers_dict) diff --git a/flytekit/tools/subprocess.py b/flytekit/tools/subprocess.py index 01ed6e2bd2..58569bf8d8 100644 --- a/flytekit/tools/subprocess.py +++ b/flytekit/tools/subprocess.py @@ -1,8 +1,9 @@ -import logging import shlex as _schlex import subprocess as _subprocess import tempfile as _tempfile +from flytekit.loggers import logger + def check_call(cmd_args, **kwargs): if not isinstance(cmd_args, list): @@ -15,12 +16,12 @@ def check_call(cmd_args, **kwargs): # Dump sub-process' std out into current std out std_out.seek(0) - logging.info("Output of command '{}':\n{}\n".format(cmd_args, std_out.read())) + logger.info("Output of command '{}':\n{}\n".format(cmd_args, std_out.read())) if ret_code != 0: std_err.seek(0) err_str = std_err.read() - logging.error("Error from command '{}':\n{}\n".format(cmd_args, err_str)) + logger.error("Error from command '{}':\n{}\n".format(cmd_args, err_str)) raise Exception( "Called process exited with error code: {}. Stderr dump:\n\n{}".format(ret_code, err_str) diff --git a/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker/training.py b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker/training.py index 557165259a..77d5b40781 100644 --- a/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker/training.py +++ b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker/training.py @@ -1,4 +1,3 @@ -import logging import typing from dataclasses import dataclass from typing import Any, Callable, Dict, TypeVar @@ -9,6 +8,7 @@ import flytekit from flytekit import ExecutionParameters, FlyteContextManager, PythonFunctionTask, kwtypes from flytekit.extend import ExecutionState, IgnoreOutputs, Interface, PythonTask, SerializationSettings, TaskPlugins +from flytekit.loggers import logger from flytekit.types.directory.types import FlyteDirectory from flytekit.types.file import FlyteFile @@ -152,7 +152,7 @@ def pre_execute(self, user_params: ExecutionParameters) -> ExecutionParameters: if the number of execution instances is > 1. Otherwise this is considered to be a single node execution """ if self._is_distributed(): - logging.info("Distributed context detected!") + logger.info("Distributed context detected!") exec_state = FlyteContextManager.current_context().execution_state if exec_state and exec_state.mode == ExecutionState.Mode.TASK_EXECUTION: """ @@ -176,10 +176,10 @@ def post_execute(self, user_params: ExecutionParameters, rval: Any) -> Any: return a None """ if self._is_distributed(): - logging.info("Distributed context detected!") + logger.info("Distributed context detected!") dctx = flytekit.current_context().distributed_training_context if not self.task_config.should_persist_output(dctx): - logging.info("output persistence predicate not met, Flytekit will ignore outputs") + logger.info("output persistence predicate not met, Flytekit will ignore outputs") raise IgnoreOutputs(f"Distributed context - Persistence predicate not met. Ignoring outputs - {dctx}") return rval diff --git a/plugins/flytekit-dolt/flytekitplugins/dolt/schema.py b/plugins/flytekit-dolt/flytekitplugins/dolt/schema.py index 6e51f079c4..8f6867b47f 100644 --- a/plugins/flytekit-dolt/flytekitplugins/dolt/schema.py +++ b/plugins/flytekit-dolt/flytekitplugins/dolt/schema.py @@ -1,4 +1,3 @@ -import logging import tempfile import typing from dataclasses import dataclass @@ -17,8 +16,6 @@ from flytekit.models.literals import Literal, Scalar from flytekit.models.types import LiteralType -logger = logging.getLogger("flytekitplugins.dolt") - @dataclass_json @dataclass diff --git a/plugins/flytekit-greatexpectations/flytekitplugins/great_expectations/schema.py b/plugins/flytekit-greatexpectations/flytekitplugins/great_expectations/schema.py index cd6f321dd2..0d3319a07f 100644 --- a/plugins/flytekit-greatexpectations/flytekitplugins/great_expectations/schema.py +++ b/plugins/flytekit-greatexpectations/flytekitplugins/great_expectations/schema.py @@ -1,5 +1,4 @@ import datetime -import logging import os import typing from dataclasses import dataclass @@ -14,6 +13,7 @@ from flytekit import FlyteContext from flytekit.extend import TypeEngine, TypeTransformer +from flytekit.loggers import logger from flytekit.models import types as _type_models from flytekit.models.literals import Literal, Primitive, Scalar from flytekit.models.types import LiteralType @@ -328,7 +328,7 @@ def to_python_value( # raise a Great Expectations' exception raise ValidationError("Validation failed!\nCOLUMN\t\tFAILED EXPECTATION\n" + result_string) - logging.info("Validation succeeded!") + logger.info("Validation succeeded!") return typing.cast(GreatExpectationsType, return_dataset) diff --git a/plugins/flytekit-greatexpectations/flytekitplugins/great_expectations/task.py b/plugins/flytekit-greatexpectations/flytekitplugins/great_expectations/task.py index 0ba60217f3..b1821e325f 100644 --- a/plugins/flytekit-greatexpectations/flytekitplugins/great_expectations/task.py +++ b/plugins/flytekit-greatexpectations/flytekitplugins/great_expectations/task.py @@ -1,5 +1,4 @@ import datetime -import logging import os import shutil from dataclasses import dataclass @@ -15,6 +14,7 @@ from flytekit import PythonInstanceTask from flytekit.core.context_manager import FlyteContext from flytekit.extend import Interface +from flytekit.loggers import logger from flytekit.types.file.file import FlyteFile from flytekit.types.schema import FlyteSchema @@ -276,6 +276,6 @@ def execute(self, **kwargs) -> Any: # raise a Great Expectations' exception raise ValidationError("Validation failed!\nCOLUMN\t\tFAILED EXPECTATION\n" + result_string) - logging.info("Validation succeeded!") + logger.info("Validation succeeded!") return final_result diff --git a/plugins/flytekit-papermill/flytekitplugins/papermill/task.py b/plugins/flytekit-papermill/flytekitplugins/papermill/task.py index 4b1e183952..a58b01d482 100644 --- a/plugins/flytekit-papermill/flytekitplugins/papermill/task.py +++ b/plugins/flytekit-papermill/flytekitplugins/papermill/task.py @@ -1,5 +1,4 @@ import json -import logging import os import typing from typing import Any @@ -13,6 +12,7 @@ from flytekit import FlyteContext, PythonInstanceTask from flytekit.core.context_manager import ExecutionParameters from flytekit.extend import Interface, TaskPlugins, TypeEngine +from flytekit.loggers import logger from flytekit.models.literals import LiteralMap from flytekit.types.file import HTMLPage, PythonNotebook @@ -200,7 +200,7 @@ def execute(self, **kwargs) -> Any: For Spark, the notebooks today need to use the new_session or just getOrCreate session and get a handle to the singleton """ - logging.info(f"Hijacking the call for task-type {self.task_type}, to call notebook.") + logger.info(f"Hijacking the call for task-type {self.task_type}, to call notebook.") # Execute Notebook via Papermill. pm.execute_notebook(self._notebook_path, self.output_notebook_path, parameters=kwargs) # type: ignore From bcfde6d1fa21c9f72f2f47e6204510d5fb926fd5 Mon Sep 17 00:00:00 2001 From: Eduardo Apolinario Date: Thu, 3 Mar 2022 14:37:23 -0800 Subject: [PATCH 105/128] Fix flytekit_compatibility/test_schema_types.py test Signed-off-by: Eduardo Apolinario Signed-off-by: maximsmol --- tests/flytekit_compatibility/test_schema_types.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/flytekit_compatibility/test_schema_types.py b/tests/flytekit_compatibility/test_schema_types.py index 6711ec34de..c151aee345 100644 --- a/tests/flytekit_compatibility/test_schema_types.py +++ b/tests/flytekit_compatibility/test_schema_types.py @@ -24,7 +24,7 @@ def test_assert_type(): schema = FlyteSchema[kwtypes(x=int, y=float)] fst = FlyteSchemaTransformer() lt = fst.get_literal_type(schema) - with pytest.raises(ValueError, match="DataFrames of type are not supported currently"): + with pytest.raises(ValueError, match="Could not convert to flyte schema"): TypeEngine.to_literal(ctx, 3, schema, lt) From f83f621bd8a7e8d2f33f0d9813d9bdbb3474dc03 Mon Sep 17 00:00:00 2001 From: Eduardo Apolinario Date: Thu, 3 Mar 2022 14:51:22 -0800 Subject: [PATCH 106/128] make lint Signed-off-by: Eduardo Apolinario --- flytekit/core/type_engine.py | 5 ++++- flytekit/types/schema/types.py | 2 +- tests/flytekit/common/parameterizers.py | 20 ++++++++++++++++++++ tests/flytekit/unit/core/test_type_engine.py | 5 ----- 4 files changed, 25 insertions(+), 7 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index c7ae0b3dd3..78e8dc6ac8 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -19,7 +19,10 @@ from google.protobuf.struct_pb2 import Struct from marshmallow_enum import EnumField, LoadDumpOptions from marshmallow_jsonschema import JSONSchema -from typing_extensions import Annotated, get_args, get_origin +from typing_extensions import Annotated +from typing_extensions import get_args +from typing_extensions import get_args as _get_args +from typing_extensions import get_origin from flytekit.core.annotation import FlyteAnnotation from flytekit.core.context_manager import FlyteContext diff --git a/flytekit/types/schema/types.py b/flytekit/types/schema/types.py index 71f1ef4ec7..8b92c901fc 100644 --- a/flytekit/types/schema/types.py +++ b/flytekit/types/schema/types.py @@ -15,7 +15,7 @@ from marshmallow import fields from flytekit.core.context_manager import FlyteContext, FlyteContextManager -from flytekit.core.type_engine import T, TypeEngine, TypeTransformer, TypeTransformerFailedError +from flytekit.core.type_engine import TypeEngine, TypeTransformer, TypeTransformerFailedError from flytekit.models.literals import Literal, Scalar, Schema from flytekit.models.types import LiteralType, SchemaType diff --git a/tests/flytekit/common/parameterizers.py b/tests/flytekit/common/parameterizers.py index 2d5504b32c..c6d9b07c82 100644 --- a/tests/flytekit/common/parameterizers.py +++ b/tests/flytekit/common/parameterizers.py @@ -257,6 +257,26 @@ ), 10, ), + ( + literals.Scalar( + union=literals.Union( + value=literals.Literal(scalar=literals.Scalar(primitive=literals.Primitive(integer=10))), + stored_type=types.LiteralType( + simple=types.SimpleType.INTEGER, structure=types.TypeStructure(tag="int") + ), + ) + ), + 10, + ), + ( + literals.Scalar( + union=literals.Union( + value=literals.Literal(scalar=literals.Scalar(primitive=literals.Primitive(string_value="test"))), + stored_type=types.LiteralType(simple=types.SimpleType.STRING, structure=types.TypeStructure(tag="str")), + ) + ), + "test", + ), ( literals.Scalar( union=literals.Union( diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index b97d77eef4..115aec6911 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -54,11 +54,6 @@ from flytekit.types.schema.types_pandas import PandasDataFrameTransformer from flytekit.types.structured.structured_dataset import StructuredDataset -try: - from typing import Annotated -except ImportError: - from typing_extensions import Annotated - T = typing.TypeVar("T") From 12cc63933662d44dc08d3e1047166def923ce1ea Mon Sep 17 00:00:00 2001 From: maximsmol Date: Mon, 7 Mar 2022 12:49:02 -0800 Subject: [PATCH 107/128] fix: annotated type conversion error Signed-off-by: maximsmol --- flytekit/core/type_engine.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 78e8dc6ac8..f70caf78fa 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -176,6 +176,9 @@ def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], exp return self._to_literal_transformer(python_val) def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[T]) -> T: + while get_origin(expected_python_type) is Annotated: + expected_python_type = _get_args(expected_python_type)[0] + if expected_python_type != self._type: raise TypeTransformerFailedError( f"Cannot convert to type {expected_python_type}, only {self._type} is supported" From 0e7569c778063543674877cba81a379e9e9077e8 Mon Sep 17 00:00:00 2001 From: maximsmol Date: Mon, 7 Mar 2022 14:06:58 -0800 Subject: [PATCH 108/128] fix: _are_types_castable based on tests Signed-off-by: maximsmol --- flytekit/core/type_engine.py | 27 ++-- flytekit/models/types.py | 2 +- .../unit/core/test_are_types_castable.py | 138 ++++++++++++++++++ 3 files changed, 152 insertions(+), 15 deletions(-) create mode 100644 tests/flytekit/unit/core/test_are_types_castable.py diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index f70caf78fa..986d0e10e3 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -877,13 +877,13 @@ def _type_essence(x: LiteralType) -> LiteralType: def _are_types_castable(upstream: LiteralType, downstream: LiteralType) -> bool: if upstream.collection_type is not None: - if upstream.collection_type is None: + if downstream.collection_type is None: return False return _are_types_castable(upstream.collection_type, downstream.collection_type) if upstream.map_value_type is not None: - if upstream.map_value_type is None: + if downstream.map_value_type is None: return False return _are_types_castable(upstream.map_value_type, downstream.map_value_type) @@ -919,19 +919,18 @@ def _are_types_castable(upstream: LiteralType, downstream: LiteralType) -> bool: return True - if downstream.union_type is not None: - if upstream.union_type is not None: - # for each upstream variant, there must be a compatible type downstream - for v in upstream.union_type: - if not _are_types_castable(v, downstream): - return False - return True + if upstream.union_type is not None: + # for each upstream variant, there must be a compatible type downstream + for v in upstream.union_type.variants: + if not _are_types_castable(v, downstream): + return False + return True - else: - # there must be a compatible downstream type - for v in downstream.union_type.variants: - if _are_types_castable(upstream, v): - return True + if downstream.union_type is not None: + # there must be a compatible downstream type + for v in downstream.union_type.variants: + if _are_types_castable(upstream, v): + return True if upstream.enum_type is not None: # enums are castable to string diff --git a/flytekit/models/types.py b/flytekit/models/types.py index c8cc1307d3..5d33a23427 100644 --- a/flytekit/models/types.py +++ b/flytekit/models/types.py @@ -258,7 +258,7 @@ def __init__( :param flytekit.models.core.types.TypeStructure structure: Type matching hints :param flytekit.models.core.types.StructuredDatasetType structured_dataset_type: structured dataset :param dict[Text, T] metadata: Additional data describing the type - :param flytekit.models.annotation.FlyteAnnotation annotation: Additional data + :param flytekit.models.annotation.TypeAnnotation annotation: Additional data describing the type _intended to be saturated by the client_ """ self._simple = simple diff --git a/tests/flytekit/unit/core/test_are_types_castable.py b/tests/flytekit/unit/core/test_are_types_castable.py new file mode 100644 index 0000000000..2f7edcdf4c --- /dev/null +++ b/tests/flytekit/unit/core/test_are_types_castable.py @@ -0,0 +1,138 @@ +from flytekit.core.type_engine import _are_types_castable +from flytekit.models.annotation import TypeAnnotation +from flytekit.models.core.types import EnumType +from flytekit.models.types import LiteralType, SimpleType, StructuredDatasetType, TypeStructure, UnionType + +str_type = LiteralType(simple=SimpleType.STRING) +int_type = LiteralType(simple=SimpleType.INTEGER) +none_type = LiteralType(simple=SimpleType.NONE) +bool_type = LiteralType(simple=SimpleType.BOOLEAN) + +str_or_int = LiteralType(union_type=UnionType([str_type, int_type])) +int_or_str = LiteralType(union_type=UnionType([int_type, str_type])) +str_or_int_or_bool = LiteralType(union_type=UnionType([str_type, int_type, bool_type])) +optional_str = LiteralType(union_type=UnionType([str_type, none_type])) + + +def test_simple(): + assert _are_types_castable(str_type, str_type) + assert not _are_types_castable(str_type, int_type) + assert not _are_types_castable(int_type, str_type) + + +def test_metadata(): + a = LiteralType(simple=SimpleType.STRING, metadata={"test": 456}) + assert _are_types_castable( + a, + LiteralType(simple=SimpleType.STRING, metadata={"test": 123}), + ) + # must not clobber metadata + assert a.metadata == {"test": 456} + + +def test_annotation(): + a = LiteralType(simple=SimpleType.STRING, annotation=TypeAnnotation(annotations={"test": 456})) + assert _are_types_castable( + a, + LiteralType(simple=SimpleType.STRING, annotation=TypeAnnotation(annotations={"test": 123})), + ) + # must not clobber annotation + assert a.annotation.annotations == {"test": 456} + + +def test_structure(): + a = LiteralType(simple=SimpleType.STRING, structure=TypeStructure(tag="a")) + assert _are_types_castable( + a, + LiteralType(simple=SimpleType.STRING, structure=TypeStructure(tag="b")), + ) + # must not clobber annotation + assert a.structure.tag == "a" + + +def test_non_nullable(): + assert not _are_types_castable(none_type, str_type) + assert not _are_types_castable(str_type, none_type) + assert _are_types_castable(none_type, none_type) + + +def test_collection(): + assert _are_types_castable(LiteralType(collection_type=str_type), LiteralType(collection_type=str_type)) + assert not _are_types_castable(LiteralType(collection_type=str_type), LiteralType(collection_type=int_type)) + + +def test_map(): + assert _are_types_castable(LiteralType(map_value_type=str_type), LiteralType(map_value_type=str_type)) + assert not _are_types_castable(LiteralType(map_value_type=str_type), LiteralType(map_value_type=int_type)) + + +def test_structured_dataset(): + x = LiteralType( + structured_dataset_type=StructuredDatasetType( + columns=[ + StructuredDatasetType.DatasetColumn("a", str_type), + StructuredDatasetType.DatasetColumn("b", int_type), + ], + format="abc", + external_schema_type="zzz", + external_schema_bytes=b"zzz", + ) + ) + assert _are_types_castable(x, x) + + +def test_enum(): + e = LiteralType(enum_type=EnumType(["a", "b"])) + # enum is a str + assert _are_types_castable(e, str_type) + # a str is not necessarily an enum + assert not _are_types_castable(str_type, e) + + +def test_union(): + # str can be expanded to str | int + assert _are_types_castable(str_type, str_or_int) + # str can be expanded to int | str + assert _are_types_castable(str_type, int_or_str) + # str | int cannot be narrowed to str + assert not _are_types_castable(str_or_int, str_type) + # str | int == str | int + assert _are_types_castable(str_or_int, str_or_int) + # int | str == str | int + assert _are_types_castable(int_or_str, str_or_int) + # str is Optional[str] + assert _are_types_castable(str_type, optional_str) + # None is Optional[str] + assert _are_types_castable(none_type, optional_str) + # bool is not Optional[str] + assert not _are_types_castable(bool_type, optional_str) + + +def test_collection_union(): + # a list of str is a list of (str | int) + assert _are_types_castable(LiteralType(collection_type=str_type), LiteralType(collection_type=str_or_int)) + # a list of int is a list of (str | int) + assert _are_types_castable(LiteralType(collection_type=int_type), LiteralType(collection_type=str_or_int)) + # a list of str or a list of int is a list of (str | int) + assert _are_types_castable( + LiteralType( + union_type=UnionType([LiteralType(collection_type=int_type), LiteralType(collection_type=str_type)]) + ), + LiteralType(collection_type=str_or_int), + ) + assert _are_types_castable( + LiteralType( + union_type=UnionType([LiteralType(collection_type=int_type), LiteralType(collection_type=str_type)]) + ), + LiteralType(collection_type=str_or_int_or_bool), + ) + # a list of str or a list of bool is not a list of (str | int) + assert not _are_types_castable( + LiteralType( + union_type=UnionType([LiteralType(collection_type=int_type), LiteralType(collection_type=bool_type)]) + ), + LiteralType(collection_type=str_or_int), + ) + # not the other way around + assert not _are_types_castable(LiteralType(collection_type=str_or_int), LiteralType(collection_type=str_type)) + assert not _are_types_castable(LiteralType(collection_type=str_or_int), LiteralType(collection_type=int_type)) From 708089848161731af92425380121320c5564e6d9 Mon Sep 17 00:00:00 2001 From: maximsmol Date: Mon, 7 Mar 2022 17:31:19 -0800 Subject: [PATCH 109/128] fix: test failing if using random order Signed-off-by: maximsmol --- flytekit/core/type_engine.py | 2 +- tests/flytekit/unit/core/test_type_engine.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 986d0e10e3..b2b44501ca 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -168,7 +168,7 @@ def __init__( self._from_literal_transformer = from_literal_transformer def get_literal_type(self, t: Type[T] = None) -> LiteralType: - return self._lt + return LiteralType.from_flyte_idl(self._lt.to_flyte_idl()) def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], expected: LiteralType) -> Literal: if type(python_val) != self._type: diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 115aec6911..c22ee2c475 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -897,11 +897,11 @@ def test_union_of_lists(): lt = TypeEngine.to_literal_type(pt) assert lt.union_type.variants == [ LiteralType( - collection_type=LiteralType(simple=SimpleType.INTEGER, structure=TypeStructure(tag="int")), + collection_type=LiteralType(simple=SimpleType.INTEGER), structure=TypeStructure(tag="Typed List"), ), LiteralType( - collection_type=LiteralType(simple=SimpleType.STRING, structure=TypeStructure(tag="str")), + collection_type=LiteralType(simple=SimpleType.STRING), structure=TypeStructure(tag="Typed List"), ), ] From 4c925a1a30f4ae2d8ef01b0306323687d829f3df Mon Sep 17 00:00:00 2001 From: maximsmol Date: Tue, 8 Mar 2022 09:43:02 -0800 Subject: [PATCH 110/128] Merge branch 'master' into maximsmol/union_type Signed-off-by: maximsmol --- .github/workflows/pythonbuild.yml | 1 - dev-requirements.txt | 7 +- doc-requirements.txt | 2 +- flytekit/clis/sdk_in_container/pyflyte.py | 8 --- flytekit/clis/sdk_in_container/serialize.py | 6 -- flytekit/core/tracker.py | 2 +- flytekit/core/type_engine.py | 7 +- flytekit/models/types.py | 8 +-- flytekit/types/file/file.py | 7 +- requirements-spark2.txt | 2 +- requirements.txt | 10 +-- tests/flytekit/common/parameterizers.py | 68 ------------------- .../workflows/requirements.txt | 2 +- tests/flytekit/unit/core/test_type_engine.py | 9 ++- .../test_schema_types.py | 2 +- 15 files changed, 28 insertions(+), 113 deletions(-) diff --git a/.github/workflows/pythonbuild.yml b/.github/workflows/pythonbuild.yml index 5af4b5f634..681d923f87 100644 --- a/.github/workflows/pythonbuild.yml +++ b/.github/workflows/pythonbuild.yml @@ -106,7 +106,6 @@ jobs: pip install -e . if [ -f dev-requirements.txt ]; then pip install -r dev-requirements.txt; fi pip install --no-deps -U https://github.com/flyteorg/flytekit/archive/${{ github.sha }}.zip#egg=flytekit - pip install --no-deps -U "git+https://github.com/maximsmol/flyteidl.git@maximsmol/union-types#egg=flyteidl" pip freeze - name: Test with coverage run: | diff --git a/dev-requirements.txt b/dev-requirements.txt index af1217527e..f8bd6af672 100644 --- a/dev-requirements.txt +++ b/dev-requirements.txt @@ -129,7 +129,7 @@ google-auth==2.6.0 # via # google-api-core # google-cloud-core -google-cloud-bigquery==2.34.1 +google-cloud-bigquery==2.34.0 # via -r dev-requirements.in google-cloud-bigquery-storage==2.12.0 # via -r dev-requirements.in @@ -137,7 +137,7 @@ google-cloud-core==2.2.2 # via google-cloud-bigquery google-crc32c==1.3.0 # via google-resumable-media -google-resumable-media==2.3.1 +google-resumable-media==2.3.0 # via google-cloud-bigquery googleapis-common-protos==1.55.0 # via @@ -256,6 +256,7 @@ proto-plus==1.20.3 protobuf==3.19.4 # via # -c requirements.txt + # flyteidl # flytekit # google-api-core # google-cloud-bigquery @@ -335,7 +336,7 @@ pyyaml==5.4.1 # -c requirements.txt # docker-compose # pre-commit -regex==2022.3.2 +regex==2022.1.18 # via # -c requirements.txt # docker-image-py diff --git a/doc-requirements.txt b/doc-requirements.txt index c4c496447d..07ef3d0d5b 100644 --- a/doc-requirements.txt +++ b/doc-requirements.txt @@ -165,7 +165,7 @@ pytz==2021.3 # pandas pyyaml==6.0 # via sphinx-autoapi -regex==2022.3.2 +regex==2022.1.18 # via docker-image-py requests==2.27.1 # via diff --git a/flytekit/clis/sdk_in_container/pyflyte.py b/flytekit/clis/sdk_in_container/pyflyte.py index 56b8328701..f370c82942 100644 --- a/flytekit/clis/sdk_in_container/pyflyte.py +++ b/flytekit/clis/sdk_in_container/pyflyte.py @@ -8,14 +8,11 @@ from flytekit.clis.sdk_in_container.local_cache import local_cache from flytekit.clis.sdk_in_container.package import package from flytekit.clis.sdk_in_container.serialize import serialize -from flytekit.configuration import internal as _internal_config from flytekit.configuration import platform as _platform_config -from flytekit.configuration import sdk as _sdk_config from flytekit.configuration import set_flyte_config_file from flytekit.configuration.internal import CONFIGURATION_PATH from flytekit.configuration.platform import URL as _URL from flytekit.configuration.sdk import WORKFLOW_PACKAGES as _WORKFLOW_PACKAGES -from flytekit.loggers import cli_logger def validate_package(ctx, param, values): @@ -58,11 +55,6 @@ def main(ctx, config=None, pkgs=None, insecure=None): """ update_configuration_file(config) - # Update the logger if it's set - log_level = _internal_config.LOGGING_LEVEL.get() or _sdk_config.LOGGING_LEVEL.get() - if log_level is not None: - cli_logger.getLogger().setLevel(log_level) - ctx.obj = dict() # Determine SSL. Note that the insecure option in this command is not a flag because we want to look diff --git a/flytekit/clis/sdk_in_container/serialize.py b/flytekit/clis/sdk_in_container/serialize.py index 4668b34b1d..09c06f3d0f 100644 --- a/flytekit/clis/sdk_in_container/serialize.py +++ b/flytekit/clis/sdk_in_container/serialize.py @@ -1,4 +1,3 @@ -import logging import math as _math import os as _os import sys @@ -21,7 +20,6 @@ from flytekit.core.workflow import WorkflowBase from flytekit.exceptions.scopes import system_entry_point from flytekit.exceptions.user import FlyteValidationException -from flytekit.loggers import cli_logger from flytekit.models import launch_plan as _launch_plan_models from flytekit.models import task as task_models from flytekit.models.admin import workflow as admin_workflow_models @@ -295,8 +293,6 @@ def serialize(ctx, image, local_source_root, in_container_config_path, in_contai @click.option("-f", "--folder", type=click.Path(exists=True)) @click.pass_context def workflows(ctx, folder=None): - cli_logger.getLogger().setLevel(logging.DEBUG) - if folder: click.echo(f"Writing output to {folder}") @@ -323,8 +319,6 @@ def fast(ctx): @click.option("-f", "--folder", type=click.Path(exists=True)) @click.pass_context def fast_workflows(ctx, folder=None): - cli_logger.getLogger().setLevel(logging.DEBUG) - if folder: click.echo(f"Writing output to {folder}") diff --git a/flytekit/core/tracker.py b/flytekit/core/tracker.py index d72adb6b6a..1cd9306f44 100644 --- a/flytekit/core/tracker.py +++ b/flytekit/core/tracker.py @@ -21,7 +21,7 @@ class InstanceTrackingMeta(type): def _find_instance_module(): frame = _inspect.currentframe() while frame: - if frame.f_code.co_name == "": + if frame.f_code.co_name == "" and "__name__" in frame.f_globals: return frame.f_globals["__name__"] frame = frame.f_back return None diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index b2b44501ca..28616679b4 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -19,10 +19,6 @@ from google.protobuf.struct_pb2 import Struct from marshmallow_enum import EnumField, LoadDumpOptions from marshmallow_jsonschema import JSONSchema -from typing_extensions import Annotated -from typing_extensions import get_args -from typing_extensions import get_args as _get_args -from typing_extensions import get_origin from flytekit.core.annotation import FlyteAnnotation from flytekit.core.context_manager import FlyteContext @@ -54,7 +50,6 @@ except ImportError: from typing_extensions import Annotated, get_args, get_origin - T = typing.TypeVar("T") DEFINITIONS = "definitions" @@ -177,7 +172,7 @@ def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], exp def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[T]) -> T: while get_origin(expected_python_type) is Annotated: - expected_python_type = _get_args(expected_python_type)[0] + expected_python_type = get_args(expected_python_type)[0] if expected_python_type != self._type: raise TypeTransformerFailedError( diff --git a/flytekit/models/types.py b/flytekit/models/types.py index 5d33a23427..3c58c46b61 100644 --- a/flytekit/models/types.py +++ b/flytekit/models/types.py @@ -239,9 +239,9 @@ def __init__( blob=None, enum_type=None, union_type=None, - structure=None, structured_dataset_type=None, metadata=None, + structure=None, annotation=None, ): """ @@ -268,9 +268,9 @@ def __init__( self._blob = blob self._enum_type = enum_type self._union_type = union_type - self._structure = structure self._structured_dataset_type = structured_dataset_type self._metadata = metadata + self._structure = structure self._annotation = annotation @property @@ -354,12 +354,12 @@ def to_flyte_idl(self): blob=self.blob.to_flyte_idl() if self.blob is not None else None, enum_type=self.enum_type.to_flyte_idl() if self.enum_type else None, union_type=self.union_type.to_flyte_idl() if self.union_type else None, - structure=self.structure.to_flyte_idl() if self.structure else None, structured_dataset_type=self.structured_dataset_type.to_flyte_idl() if self.structured_dataset_type else None, metadata=metadata, annotation=self.annotation.to_flyte_idl() if self.annotation else None, + structure=self.structure.to_flyte_idl() if self.structure else None, ) return t @@ -383,11 +383,11 @@ def from_flyte_idl(cls, proto): blob=_core_types.BlobType.from_flyte_idl(proto.blob) if proto.HasField("blob") else None, enum_type=_core_types.EnumType.from_flyte_idl(proto.enum_type) if proto.HasField("enum_type") else None, union_type=UnionType.from_flyte_idl(proto.union_type) if proto.HasField("union_type") else None, - structure=TypeStructure.from_flyte_idl(proto.structure) if proto.HasField("structure") else None, structured_dataset_type=StructuredDatasetType.from_flyte_idl(proto.structured_dataset_type) if proto.HasField("structured_dataset_type") else None, metadata=_json_format.MessageToDict(proto.metadata) or None, + structure=TypeStructure.from_flyte_idl(proto.structure) if proto.HasField("structure") else None, annotation=TypeAnnotationModel.from_flyte_idl(proto.annotation) if proto.HasField("annotation") else None, ) diff --git a/flytekit/types/file/file.py b/flytekit/types/file/file.py index a215b6b851..744f56de6b 100644 --- a/flytekit/types/file/file.py +++ b/flytekit/types/file/file.py @@ -38,14 +38,17 @@ class FlyteFile(os.PathLike, typing.Generic[T]): the contents of the files to the blob store connected with your Flyte installation. That is, the Python native literal that represents a file is typically just the path to the file on the local filesystem. However in Flyte, an instance of a file is represented by a :py:class:`Blob ` literal, - with the ``uri`` field set to the location in the Flyte blob store (AWS/GCS etc.). + with the ``uri`` field set to the location in the Flyte blob store (AWS/GCS etc.). Take a look at the + :std:ref:`data handling doc ` for a deeper discussion. We decided to not support ``pathlib.Path`` as an input/output type because if you wanted the automatic upload/download behavior, you should just use the ``FlyteFile`` type. If you do not, then a ``str`` works just as well. The prefix for where uploads go is set by the raw output data prefix setting, which should be set at registration - time. See the flytectl option for more information. + time in the launch plan. See the option listed under ``flytectl register examples --help`` for more information. + If not set in the launch plan, then your Flyte backend will specify a default. This default is itself configurable + as well. Contact your Flyte platform administrators to change or ascertain the value. In short, if a task returns ``"/path/to/file"`` and the task's signature is set to return ``FlyteFile``, then the contents of ``/path/to/file`` are uploaded. diff --git a/requirements-spark2.txt b/requirements-spark2.txt index bd072c15d6..3ad6b4d78e 100644 --- a/requirements-spark2.txt +++ b/requirements-spark2.txt @@ -136,7 +136,7 @@ pytz==2021.3 # pandas pyyaml==5.4.1 # via -r requirements.in -regex==2022.3.2 +regex==2022.1.18 # via docker-image-py requests==2.27.1 # via diff --git a/requirements.txt b/requirements.txt index 0161dded53..d1c1d6122b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -124,15 +124,7 @@ python-dateutil==2.8.2 # pandas python-json-logger==2.0.2 # via flytekit -<<<<<<< HEAD -<<<<<<< HEAD python-slugify==6.1.1 -======= -python-slugify==6.1.0 ->>>>>>> 6a1ceea5 (Bump idl (#862)) -======= -python-slugify==6.1.1 ->>>>>>> 0da523c7 (Caching of offloaded objects (#762)) # via cookiecutter pytimeparse==1.1.8 # via flytekit @@ -142,7 +134,7 @@ pytz==2021.3 # pandas pyyaml==5.4.1 # via -r requirements.in -regex==2022.3.2 +regex==2022.1.18 # via docker-image-py requests==2.27.1 # via diff --git a/tests/flytekit/common/parameterizers.py b/tests/flytekit/common/parameterizers.py index c6d9b07c82..4b48fbcfb9 100644 --- a/tests/flytekit/common/parameterizers.py +++ b/tests/flytekit/common/parameterizers.py @@ -178,74 +178,6 @@ timedelta(seconds=5), ), (literals.Scalar(none_type=literals.Void()), None), - ( - literals.Scalar( - blob=literals.Blob( - literals.BlobMetadata(_core_types.BlobType("csv", _core_types.BlobType.BlobDimensionality.SINGLE)), - "s3://some/where", - ) - ), - _blob_impl.Blob("s3://some/where", format="csv"), - ), - ( - literals.Scalar( - blob=literals.Blob( - literals.BlobMetadata(_core_types.BlobType("", _core_types.BlobType.BlobDimensionality.SINGLE)), - "s3://some/where", - ) - ), - _blob_impl.Blob("s3://some/where"), - ), - ( - literals.Scalar( - blob=literals.Blob( - literals.BlobMetadata(_core_types.BlobType("csv", _core_types.BlobType.BlobDimensionality.MULTIPART)), - "s3://some/where/", - ) - ), - _blob_impl.MultiPartBlob("s3://some/where/", format="csv"), - ), - ( - literals.Scalar( - blob=literals.Blob( - literals.BlobMetadata(_core_types.BlobType("", _core_types.BlobType.BlobDimensionality.MULTIPART)), - "s3://some/where/", - ) - ), - _blob_impl.MultiPartBlob("s3://some/where/"), - ), - ( - literals.Scalar( - schema=literals.Schema( - "s3://some/where/", - types.SchemaType( - [ - types.SchemaType.SchemaColumn("a", types.SchemaType.SchemaColumn.SchemaColumnType.INTEGER), - types.SchemaType.SchemaColumn("b", types.SchemaType.SchemaColumn.SchemaColumnType.BOOLEAN), - types.SchemaType.SchemaColumn("c", types.SchemaType.SchemaColumn.SchemaColumnType.DATETIME), - types.SchemaType.SchemaColumn("d", types.SchemaType.SchemaColumn.SchemaColumnType.DURATION), - types.SchemaType.SchemaColumn("e", types.SchemaType.SchemaColumn.SchemaColumnType.FLOAT), - types.SchemaType.SchemaColumn("f", types.SchemaType.SchemaColumn.SchemaColumnType.STRING), - ] - ), - ) - ), - _schema_impl.Schema( - "s3://some/where/", - _schema_impl.SchemaType.promote_from_model( - types.SchemaType( - [ - types.SchemaType.SchemaColumn("a", types.SchemaType.SchemaColumn.SchemaColumnType.INTEGER), - types.SchemaType.SchemaColumn("b", types.SchemaType.SchemaColumn.SchemaColumnType.BOOLEAN), - types.SchemaType.SchemaColumn("c", types.SchemaType.SchemaColumn.SchemaColumnType.DATETIME), - types.SchemaType.SchemaColumn("d", types.SchemaType.SchemaColumn.SchemaColumnType.DURATION), - types.SchemaType.SchemaColumn("e", types.SchemaType.SchemaColumn.SchemaColumnType.FLOAT), - types.SchemaType.SchemaColumn("f", types.SchemaType.SchemaColumn.SchemaColumnType.STRING), - ] - ) - ), - ), - ), ( literals.Scalar( union=literals.Union( diff --git a/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt b/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt index fb365ed51d..15a1f89700 100644 --- a/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt +++ b/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt @@ -142,7 +142,7 @@ pytz==2021.3 # via # flytekit # pandas -regex==2022.3.2 +regex==2022.1.18 # via docker-image-py requests==2.27.1 # via diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index c22ee2c475..8c1970a5ed 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -19,7 +19,6 @@ from pandas._testing import assert_frame_equal from typing_extensions import Annotated -import flytekit.common.exceptions.user as user_exceptions from flytekit import kwtypes from flytekit.core.annotation import FlyteAnnotation from flytekit.core.context_manager import FlyteContext, FlyteContextManager @@ -382,6 +381,14 @@ def test_guessing_basic(): pt = TypeEngine.guess_python_type(lt) assert pt is FlytePickle + lt = model_types.LiteralType( + blob=BlobType( + format=FlytePickleTransformer.PYTHON_PICKLE_FORMAT, dimensionality=BlobType.BlobDimensionality.SINGLE + ) + ) + pt = TypeEngine.guess_python_type(lt) + assert pt is FlytePickle + def test_guessing_containers(): b = model_types.LiteralType(simple=model_types.SimpleType.BOOLEAN) diff --git a/tests/flytekit_compatibility/test_schema_types.py b/tests/flytekit_compatibility/test_schema_types.py index c151aee345..6711ec34de 100644 --- a/tests/flytekit_compatibility/test_schema_types.py +++ b/tests/flytekit_compatibility/test_schema_types.py @@ -24,7 +24,7 @@ def test_assert_type(): schema = FlyteSchema[kwtypes(x=int, y=float)] fst = FlyteSchemaTransformer() lt = fst.get_literal_type(schema) - with pytest.raises(ValueError, match="Could not convert to flyte schema"): + with pytest.raises(ValueError, match="DataFrames of type are not supported currently"): TypeEngine.to_literal(ctx, 3, schema, lt) From b4526090851ed66f36f07a6a05ad3ebe72c62cf8 Mon Sep 17 00:00:00 2001 From: maximsmol Date: Tue, 8 Mar 2022 09:45:51 -0800 Subject: [PATCH 111/128] fix: merge issue Signed-off-by: maximsmol --- flytekit/models/types.py | 1 + 1 file changed, 1 insertion(+) diff --git a/flytekit/models/types.py b/flytekit/models/types.py index 3c58c46b61..6a5d5b7e29 100644 --- a/flytekit/models/types.py +++ b/flytekit/models/types.py @@ -311,6 +311,7 @@ def union_type(self) -> UnionType: def structure(self) -> TypeStructure: return self._structure + @property def structured_dataset_type(self) -> StructuredDatasetType: return self._structured_dataset_type From adcbb79fbaec53438dae478d6492d0f67a9dd935 Mon Sep 17 00:00:00 2001 From: maximsmol Date: Tue, 8 Mar 2022 10:03:47 -0800 Subject: [PATCH 112/128] fix: requirements Signed-off-by: maximsmol --- dev-requirements.txt | 44 ++++--------------- doc-requirements.txt | 28 +++--------- requirements-spark2.txt | 18 +------- requirements.txt | 18 +------- .../workflows/requirements.txt | 19 ++------ 5 files changed, 21 insertions(+), 106 deletions(-) diff --git a/dev-requirements.txt b/dev-requirements.txt index f8bd6af672..2dc8af00cc 100644 --- a/dev-requirements.txt +++ b/dev-requirements.txt @@ -32,7 +32,6 @@ certifi==2021.10.8 # requests cffi==1.15.0 # via - # -c requirements.txt # bcrypt # cryptography # pynacl @@ -72,10 +71,7 @@ croniter==1.3.4 # -c requirements.txt # flytekit cryptography==36.0.1 - # via - # -c requirements.txt - # paramiko - # secretstorage + # via paramiko dataclasses-json==0.5.6 # via # -c requirements.txt @@ -116,11 +112,7 @@ docstring-parser==0.13 # flytekit filelock==3.6.0 # via virtualenv -flyteidl==0.23.0 - # via - # -c requirements.txt - # flytekit -google-api-core[grpc]==2.5.0 +google-api-core[grpc]==2.6.1 # via # google-cloud-bigquery # google-cloud-bigquery-storage @@ -129,20 +121,19 @@ google-auth==2.6.0 # via # google-api-core # google-cloud-core -google-cloud-bigquery==2.34.0 +google-cloud-bigquery==2.34.2 # via -r dev-requirements.in google-cloud-bigquery-storage==2.12.0 # via -r dev-requirements.in -google-cloud-core==2.2.2 +google-cloud-core==2.2.3 # via google-cloud-bigquery google-crc32c==1.3.0 # via google-resumable-media -google-resumable-media==2.3.0 +google-resumable-media==2.3.1 # via google-cloud-bigquery googleapis-common-protos==1.55.0 # via # -c requirements.txt - # flyteidl # google-api-core # grpcio-status grpcio==1.44.0 @@ -166,11 +157,6 @@ importlib-metadata==4.11.2 # keyring iniconfig==1.1.1 # via pytest -jeepney==0.7.1 - # via - # -c requirements.txt - # keyring - # secretstorage jinja2==3.0.3 # via # -c requirements.txt @@ -256,18 +242,12 @@ proto-plus==1.20.3 protobuf==3.19.4 # via # -c requirements.txt - # flyteidl # flytekit # google-api-core # google-cloud-bigquery # googleapis-common-protos # grpcio-status # proto-plus - # protoc-gen-swagger -protoc-gen-swagger==0.1.0 - # via - # -c requirements.txt - # flyteidl py==1.11.0 # via # -c requirements.txt @@ -284,9 +264,7 @@ pyasn1==0.4.8 pyasn1-modules==0.2.8 # via google-auth pycparser==2.21 - # via - # -c requirements.txt - # cffi + # via cffi pynacl==1.5.0 # via paramiko pyparsing==3.0.7 @@ -336,7 +314,7 @@ pyyaml==5.4.1 # -c requirements.txt # docker-compose # pre-commit -regex==2022.1.18 +regex==2022.3.2 # via # -c requirements.txt # docker-image-py @@ -350,7 +328,7 @@ requests==2.27.1 # google-api-core # google-cloud-bigquery # responses -responses==0.18.0 +responses==0.19.0 # via # -c requirements.txt # flytekit @@ -360,10 +338,6 @@ retry==0.9.2 # flytekit rsa==4.8 # via google-auth -secretstorage==3.3.1 - # via - # -c requirements.txt - # keyring six==1.16.0 # via # -c requirements.txt @@ -414,7 +388,7 @@ urllib3==1.26.8 # flytekit # requests # responses -virtualenv==20.13.2 +virtualenv==20.13.3 # via pre-commit websocket-client==0.59.0 # via diff --git a/doc-requirements.txt b/doc-requirements.txt index 07ef3d0d5b..662cd31900 100644 --- a/doc-requirements.txt +++ b/doc-requirements.txt @@ -42,9 +42,7 @@ cookiecutter==1.7.3 croniter==1.3.4 # via flytekit cryptography==36.0.1 - # via - # -r doc-requirements.in - # secretstorage + # via -r doc-requirements.in css-html-js-minify==2.5.5 # via sphinx-material dataclasses-json==0.5.6 @@ -63,12 +61,8 @@ docutils==0.17.1 # via # sphinx # sphinx-panels -flyteidl==0.23.0 - # via flytekit furo @ git+https://github.com/flyteorg/furo@main # via -r doc-requirements.in -googleapis-common-protos==1.55.0 - # via flyteidl grpcio==1.44.0 # via # -r doc-requirements.in @@ -81,10 +75,6 @@ importlib-metadata==4.11.2 # via # keyring # sphinx -jeepney==0.7.1 - # via - # keyring - # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -114,7 +104,7 @@ mypy-extensions==0.4.3 # via typing-inspect natsort==8.1.0 # via flytekit -numpy==1.22.2 +numpy==1.22.3 # via # pandas # pyarrow @@ -125,13 +115,7 @@ pandas==1.4.1 poyo==0.5.0 # via cookiecutter protobuf==3.19.4 - # via - # flyteidl - # flytekit - # googleapis-common-protos - # protoc-gen-swagger -protoc-gen-swagger==0.1.0 - # via flyteidl + # via flytekit py==1.11.0 # via retry pyarrow==6.0.1 @@ -165,7 +149,7 @@ pytz==2021.3 # pandas pyyaml==6.0 # via sphinx-autoapi -regex==2022.1.18 +regex==2022.3.2 # via docker-image-py requests==2.27.1 # via @@ -173,12 +157,10 @@ requests==2.27.1 # flytekit # responses # sphinx -responses==0.18.0 +responses==0.19.0 # via flytekit retry==0.9.2 # via flytekit -secretstorage==3.3.1 - # via keyring six==1.16.0 # via # cookiecutter diff --git a/requirements-spark2.txt b/requirements-spark2.txt index 3ad6b4d78e..f9396f48fb 100644 --- a/requirements-spark2.txt +++ b/requirements-spark2.txt @@ -20,8 +20,6 @@ binaryornot==0.4.4 # via cookiecutter certifi==2021.10.8 # via requests -cffi==1.15.0 - # via cryptography chardet==4.0.0 # via binaryornot charset-normalizer==2.0.12 @@ -38,8 +36,6 @@ cookiecutter==1.7.3 # via flytekit croniter==1.3.4 # via flytekit -cryptography==36.0.1 - # via secretstorage dataclasses-json==0.5.6 # via flytekit decorator==5.1.1 @@ -52,8 +48,6 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.23.0 - # via flytekit googleapis-common-protos==1.55.0 # via flyteidl grpcio==1.44.0 @@ -62,10 +56,6 @@ idna==3.3 # via requests importlib-metadata==4.11.2 # via keyring -jeepney==0.7.1 - # via - # keyring - # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -114,8 +104,6 @@ py==1.11.0 # via retry pyarrow==6.0.1 # via flytekit -pycparser==2.21 - # via cffi pyrsistent==0.18.1 # via jsonschema python-dateutil==2.8.2 @@ -136,19 +124,17 @@ pytz==2021.3 # pandas pyyaml==5.4.1 # via -r requirements.in -regex==2022.1.18 +regex==2022.3.2 # via docker-image-py requests==2.27.1 # via # cookiecutter # flytekit # responses -responses==0.18.0 +responses==0.19.0 # via flytekit retry==0.9.2 # via flytekit -secretstorage==3.3.1 - # via keyring six==1.16.0 # via # cookiecutter diff --git a/requirements.txt b/requirements.txt index d1c1d6122b..9ef82cbbe6 100644 --- a/requirements.txt +++ b/requirements.txt @@ -18,8 +18,6 @@ binaryornot==0.4.4 # via cookiecutter certifi==2021.10.8 # via requests -cffi==1.15.0 - # via cryptography chardet==4.0.0 # via binaryornot charset-normalizer==2.0.12 @@ -36,8 +34,6 @@ cookiecutter==1.7.3 # via flytekit croniter==1.3.4 # via flytekit -cryptography==36.0.1 - # via secretstorage dataclasses-json==0.5.6 # via flytekit decorator==5.1.1 @@ -50,8 +46,6 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.23.0 - # via flytekit googleapis-common-protos==1.55.0 # via flyteidl grpcio==1.44.0 @@ -60,10 +54,6 @@ idna==3.3 # via requests importlib-metadata==4.11.2 # via keyring -jeepney==0.7.1 - # via - # keyring - # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -112,8 +102,6 @@ py==1.11.0 # via retry pyarrow==6.0.1 # via flytekit -pycparser==2.21 - # via cffi pyrsistent==0.18.1 # via jsonschema python-dateutil==2.8.2 @@ -134,19 +122,17 @@ pytz==2021.3 # pandas pyyaml==5.4.1 # via -r requirements.in -regex==2022.1.18 +regex==2022.3.2 # via docker-image-py requests==2.27.1 # via # cookiecutter # flytekit # responses -responses==0.18.0 +responses==0.19.0 # via flytekit retry==0.9.2 # via flytekit -secretstorage==3.3.1 - # via keyring six==1.16.0 # via # cookiecutter diff --git a/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt b/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt index 15a1f89700..72b8c22d1b 100644 --- a/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt +++ b/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt @@ -10,8 +10,6 @@ binaryornot==0.4.4 # via cookiecutter certifi==2021.10.8 # via requests -cffi==1.15.0 - # via cryptography chardet==4.0.0 # via binaryornot charset-normalizer==2.0.12 @@ -28,8 +26,6 @@ cookiecutter==1.7.3 # via flytekit croniter==1.3.4 # via flytekit -cryptography==36.0.1 - # via secretstorage cycler==0.11.0 # via matplotlib dataclasses-json==0.5.6 @@ -58,10 +54,6 @@ idna==3.3 # via requests importlib-metadata==4.11.2 # via keyring -jeepney==0.7.1 - # via - # keyring - # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -91,7 +83,7 @@ mypy-extensions==0.4.3 # via typing-inspect natsort==8.1.0 # via flytekit -numpy==1.22.2 +numpy==1.22.3 # via # matplotlib # opencv-python @@ -119,15 +111,12 @@ py==1.11.0 # via retry pyarrow==6.0.1 # via flytekit -pycparser==2.21 - # via cffi pyparsing==3.0.7 # via # matplotlib # packaging python-dateutil==2.8.2 # via - # arrow # croniter # flytekit # matplotlib @@ -142,19 +131,17 @@ pytz==2021.3 # via # flytekit # pandas -regex==2022.1.18 +regex==2022.3.2 # via docker-image-py requests==2.27.1 # via # cookiecutter # flytekit # responses -responses==0.18.0 +responses==0.19.0 # via flytekit retry==0.9.2 # via flytekit -secretstorage==3.3.1 - # via keyring six==1.16.0 # via # cookiecutter From eb75bea3378463966298b5b797da006d5ad50675 Mon Sep 17 00:00:00 2001 From: maximsmol Date: Tue, 8 Mar 2022 09:43:02 -0800 Subject: [PATCH 113/128] fix: schema transformer error Signed-off-by: maximsmol --- flytekit/types/schema/types.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/flytekit/types/schema/types.py b/flytekit/types/schema/types.py index 8b92c901fc..6f01cea085 100644 --- a/flytekit/types/schema/types.py +++ b/flytekit/types/schema/types.py @@ -371,7 +371,9 @@ def to_literal( try: h = SchemaEngine.get_handler(type(python_val)) except ValueError as e: - raise TypeTransformerFailedError(f"Could not convert {type(python_val)} to flyte schema") from e + raise TypeTransformerFailedError( + f"DataFrames of type {type(python_val)} are not supported currently" + ) from e writer = schema.open(type(python_val)) writer.write(python_val) if not h.handles_remote_io: From 849520722921f45fecaaba525526286bd88bcf92 Mon Sep 17 00:00:00 2001 From: maximsmol Date: Tue, 8 Mar 2022 11:33:43 -0800 Subject: [PATCH 114/128] fix: test Signed-off-by: maximsmol --- tests/flytekit/unit/core/test_schema_types.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/flytekit/unit/core/test_schema_types.py b/tests/flytekit/unit/core/test_schema_types.py index c151aee345..6711ec34de 100644 --- a/tests/flytekit/unit/core/test_schema_types.py +++ b/tests/flytekit/unit/core/test_schema_types.py @@ -24,7 +24,7 @@ def test_assert_type(): schema = FlyteSchema[kwtypes(x=int, y=float)] fst = FlyteSchemaTransformer() lt = fst.get_literal_type(schema) - with pytest.raises(ValueError, match="Could not convert to flyte schema"): + with pytest.raises(ValueError, match="DataFrames of type are not supported currently"): TypeEngine.to_literal(ctx, 3, schema, lt) From cb7b8a1839d43b53ee5da2832f2793b1b88d296d Mon Sep 17 00:00:00 2001 From: maximsmol Date: Wed, 9 Mar 2022 08:44:34 -0800 Subject: [PATCH 115/128] fix: merge issue Signed-off-by: maximsmol --- .github/workflows/pythonbuild.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/pythonbuild.yml b/.github/workflows/pythonbuild.yml index 681d923f87..5af4b5f634 100644 --- a/.github/workflows/pythonbuild.yml +++ b/.github/workflows/pythonbuild.yml @@ -106,6 +106,7 @@ jobs: pip install -e . if [ -f dev-requirements.txt ]; then pip install -r dev-requirements.txt; fi pip install --no-deps -U https://github.com/flyteorg/flytekit/archive/${{ github.sha }}.zip#egg=flytekit + pip install --no-deps -U "git+https://github.com/maximsmol/flyteidl.git@maximsmol/union-types#egg=flyteidl" pip freeze - name: Test with coverage run: | From cc3883c31001ff76b5e332bcadd5eb5d16d9e1de Mon Sep 17 00:00:00 2001 From: maximsmol Date: Wed, 9 Mar 2022 09:03:01 -0800 Subject: [PATCH 116/128] fix: union + annotated behavior Signed-off-by: maximsmol --- flytekit/core/type_engine.py | 25 +++++---- flytekit/models/annotation.py | 5 ++ tests/flytekit/unit/core/test_type_engine.py | 53 ++++++++++++++++++++ 3 files changed, 72 insertions(+), 11 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 28616679b4..5fca8cd654 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -653,16 +653,8 @@ def to_literal_type(cls, python_type: Type) -> LiteralType: data = x.data if data is not None: idl_type_annotation = TypeAnnotationModel(annotations=data) - return LiteralType( - simple=res.simple, - schema=res.schema, - collection_type=res.collection_type, - map_value_type=res.map_value_type, - blob=res.blob, - enum_type=res.enum_type, - metadata=res.metadata, - annotation=idl_type_annotation, - ) + res = LiteralType.from_flyte_idl(res.to_flyte_idl()) + res._annotation = idl_type_annotation return res @classmethod @@ -947,14 +939,22 @@ def __init__(self): super().__init__("Typed Union", typing.Union) def get_literal_type(self, t: Type[T]) -> Optional[LiteralType]: + if get_origin(t) is Annotated: + t = get_args(t)[0] + try: trans = [(TypeEngine.get_transformer(x), x) for x in get_args(t)] - variants = [_add_tag_to_type(t.get_literal_type(x), t.name) for (t, x) in trans] + # must go through TypeEngine.to_literal_type instead of trans.get_literal_type + # to handle Annotated + variants = [_add_tag_to_type(TypeEngine.to_literal_type(x), t.name) for (t, x) in trans] return _type_models.LiteralType(union_type=UnionType(variants)) except Exception as e: raise ValueError(f"Type of Generic Union type is not supported, {e}") def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], expected: LiteralType) -> Literal: + if get_origin(python_type) is Annotated: + python_type = get_args(python_type)[0] + found_res = False res = None res_type = None @@ -978,6 +978,9 @@ def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], exp raise TypeTransformerFailedError(f"Cannot convert from {python_val} to {python_type}") def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[T]) -> Optional[typing.Any]: + if get_origin(expected_python_type) is Annotated: + expected_python_type = get_args(expected_python_type)[0] + union_tag = None union_type = None if lv.scalar is not None and lv.scalar.union is not None: diff --git a/flytekit/models/annotation.py b/flytekit/models/annotation.py index ced935c57e..bea6b1dc60 100644 --- a/flytekit/models/annotation.py +++ b/flytekit/models/annotation.py @@ -41,3 +41,8 @@ def from_flyte_idl(cls, proto): """ return cls(annotations=_json_format.MessageToDict(proto.annotations)) + + def __eq__(self, x: object) -> bool: + if not isinstance(x, self.__class__): + return False + return self.annotations == x.annotations diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 39ea27c97c..ec070d5037 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -766,6 +766,59 @@ def test_union_type(): assert v == "hello" +def test_union_type_with_annotated(): + pt = typing.Union[ + Annotated[str, FlyteAnnotation({"hello": "world"})], Annotated[int, FlyteAnnotation({"test": 123})] + ] + lt = TypeEngine.to_literal_type(pt) + assert lt.union_type.variants == [ + LiteralType( + simple=SimpleType.STRING, structure=TypeStructure(tag="str"), annotation=TypeAnnotation({"hello": "world"}) + ), + LiteralType( + simple=SimpleType.INTEGER, structure=TypeStructure(tag="int"), annotation=TypeAnnotation({"test": 123}) + ), + ] + assert union_type_tags_unique(lt) + + ctx = FlyteContextManager.current_context() + lv = TypeEngine.to_literal(ctx, 3, pt, lt) + v = TypeEngine.to_python_value(ctx, lv, pt) + assert lv.scalar.union.stored_type.structure.tag == "int" + assert lv.scalar.union.value.scalar.primitive.integer == 3 + assert v == 3 + + lv = TypeEngine.to_literal(ctx, "hello", pt, lt) + v = TypeEngine.to_python_value(ctx, lv, pt) + assert lv.scalar.union.stored_type.structure.tag == "str" + assert lv.scalar.union.value.scalar.primitive.string_value == "hello" + assert v == "hello" + + +def test_annotated_union_type(): + pt = Annotated[typing.Union[str, int], FlyteAnnotation({"hello": "world"})] + lt = TypeEngine.to_literal_type(pt) + assert lt.union_type.variants == [ + LiteralType(simple=SimpleType.STRING, structure=TypeStructure(tag="str")), + LiteralType(simple=SimpleType.INTEGER, structure=TypeStructure(tag="int")), + ] + assert lt.annotation == TypeAnnotation({"hello": "world"}) + assert union_type_tags_unique(lt) + + ctx = FlyteContextManager.current_context() + lv = TypeEngine.to_literal(ctx, 3, pt, lt) + v = TypeEngine.to_python_value(ctx, lv, pt) + assert lv.scalar.union.stored_type.structure.tag == "int" + assert lv.scalar.union.value.scalar.primitive.integer == 3 + assert v == 3 + + lv = TypeEngine.to_literal(ctx, "hello", pt, lt) + v = TypeEngine.to_python_value(ctx, lv, pt) + assert lv.scalar.union.stored_type.structure.tag == "str" + assert lv.scalar.union.value.scalar.primitive.string_value == "hello" + assert v == "hello" + + def test_optional_type(): pt = typing.Optional[int] lt = TypeEngine.to_literal_type(pt) From a7f0c182b4b0d9c91018973e9019f2a350e44e30 Mon Sep 17 00:00:00 2001 From: Eduardo Apolinario Date: Thu, 17 Mar 2022 15:13:17 -0700 Subject: [PATCH 117/128] Regenerate requirements Signed-off-by: Eduardo Apolinario --- dev-requirements.txt | 15 +++++++++------ requirements-spark2.txt | 16 +++++++++++----- requirements.in | 1 - requirements.txt | 16 +++++++++++----- 4 files changed, 31 insertions(+), 17 deletions(-) diff --git a/dev-requirements.txt b/dev-requirements.txt index e9441758f1..550deb3d4b 100644 --- a/dev-requirements.txt +++ b/dev-requirements.txt @@ -4,10 +4,6 @@ # # make dev-requirements.txt # --e git+https://github.com/maximsmol/flyteidl.git@maximsmol/union-types#egg=flyteidl - # via - # -c requirements.txt - # flytekit -e file:.#egg=flytekit # via # -c requirements.txt @@ -120,12 +116,16 @@ docstring-parser==0.13 # flytekit filelock==3.6.0 # via virtualenv +flyteidl==0.24.0 + # via + # -c requirements.txt + # flytekit google-api-core[grpc]==2.7.1 # via # google-cloud-bigquery # google-cloud-bigquery-storage # google-cloud-core -google-auth==2.6.0 +google-auth==2.6.2 # via # google-api-core # google-cloud-core @@ -141,6 +141,7 @@ google-resumable-media==2.3.2 # via google-cloud-bigquery googleapis-common-protos==1.56.0 # via + # -c requirements.txt # flyteidl # google-api-core # grpcio-status @@ -266,7 +267,9 @@ protobuf==3.19.4 # proto-plus # protoc-gen-swagger protoc-gen-swagger==0.1.0 - # via flyteidl + # via + # -c requirements.txt + # flyteidl py==1.11.0 # via # -c requirements.txt diff --git a/requirements-spark2.txt b/requirements-spark2.txt index 586d90d4d3..ee6690d1c5 100644 --- a/requirements-spark2.txt +++ b/requirements-spark2.txt @@ -4,10 +4,6 @@ # # make requirements-spark2.txt # --e git+https://github.com/maximsmol/flyteidl.git@maximsmol/union-types#egg=flyteidl - # via - # -r requirements.in - # flytekit -e file:.#egg=flytekit # via # -r requirements-spark2.in @@ -54,6 +50,10 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit +flyteidl==0.24.0 + # via flytekit +googleapis-common-protos==1.56.0 + # via flyteidl grpcio==1.44.0 # via flytekit idna==3.3 @@ -103,7 +103,13 @@ pandas==1.3.5 poyo==0.5.0 # via cookiecutter protobuf==3.19.4 - # via flytekit + # via + # flyteidl + # flytekit + # googleapis-common-protos + # protoc-gen-swagger +protoc-gen-swagger==0.1.0 + # via flyteidl py==1.11.0 # via retry pyarrow==6.0.1 diff --git a/requirements.in b/requirements.in index 545ef2c706..0cf01a6f13 100644 --- a/requirements.in +++ b/requirements.in @@ -1,5 +1,4 @@ .[all] --e git+https://github.com/maximsmol/flyteidl.git@maximsmol/union-types#egg=flyteidl -e file:.#egg=flytekit attrs<21 # We need to restrict constrain the versions of both jsonschema and pyyaml because of docker-compose (which is diff --git a/requirements.txt b/requirements.txt index 8603b25c7c..739fd6efb6 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,10 +4,6 @@ # # make requirements.txt # --e git+https://github.com/maximsmol/flyteidl.git@maximsmol/union-types#egg=flyteidl - # via - # -r requirements.in - # flytekit -e file:.#egg=flytekit # via -r requirements.in arrow==1.2.2 @@ -52,6 +48,10 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit +flyteidl==0.24.0 + # via flytekit +googleapis-common-protos==1.56.0 + # via flyteidl grpcio==1.44.0 # via flytekit idna==3.3 @@ -101,7 +101,13 @@ pandas==1.3.5 poyo==0.5.0 # via cookiecutter protobuf==3.19.4 - # via flytekit + # via + # flyteidl + # flytekit + # googleapis-common-protos + # protoc-gen-swagger +protoc-gen-swagger==0.1.0 + # via flyteidl py==1.11.0 # via retry pyarrow==6.0.1 From 6526999bdc13cf7253572b94f4bb6621696d1d78 Mon Sep 17 00:00:00 2001 From: Eduardo Apolinario Date: Thu, 17 Mar 2022 16:51:03 -0700 Subject: [PATCH 118/128] Bump requirements in plugins Signed-off-by: Eduardo Apolinario --- plugins/flytekit-aws-athena/requirements.txt | 30 ++--- plugins/flytekit-aws-batch/requirements.txt | 30 ++--- .../flytekit-aws-sagemaker/requirements.txt | 45 ++++---- plugins/flytekit-bigquery/requirements.txt | 40 +++---- plugins/flytekit-data-fsspec/requirements.txt | 34 +++--- plugins/flytekit-dolt/requirements.txt | 30 ++--- .../requirements.txt | 73 +++++++----- plugins/flytekit-hive/requirements.txt | 30 ++--- plugins/flytekit-k8s-pod/requirements.txt | 32 +++--- plugins/flytekit-kf-mpi/requirements.txt | 30 ++--- plugins/flytekit-kf-pytorch/requirements.txt | 30 ++--- .../flytekit-kf-tensorflow/requirements.txt | 30 ++--- plugins/flytekit-modin/requirements.txt | 65 +++++++---- plugins/flytekit-pandera/requirements.txt | 30 ++--- plugins/flytekit-papermill/requirements.txt | 107 +++++++++--------- plugins/flytekit-snowflake/requirements.txt | 30 ++--- plugins/flytekit-spark/requirements.txt | 30 ++--- plugins/flytekit-sqlalchemy/requirements.txt | 32 +++--- 18 files changed, 412 insertions(+), 316 deletions(-) diff --git a/plugins/flytekit-aws-athena/requirements.txt b/plugins/flytekit-aws-athena/requirements.txt index bb210d8c13..1af8335a2c 100644 --- a/plugins/flytekit-aws-athena/requirements.txt +++ b/plugins/flytekit-aws-athena/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.10 +# This file is autogenerated by pip-compile with python 3.8 # To update, run: # # pip-compile requirements.in @@ -30,7 +30,7 @@ cookiecutter==1.7.3 # via flytekit croniter==1.3.4 # via flytekit -cryptography==36.0.1 +cryptography==36.0.2 # via secretstorage dataclasses-json==0.5.6 # via flytekit @@ -44,17 +44,17 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.23.0 +flyteidl==0.24.0 # via flytekit -flytekit==0.30.3 +flytekit==0.31.0 # via flytekitplugins-athena -googleapis-common-protos==1.55.0 +googleapis-common-protos==1.56.0 # via flyteidl grpcio==1.44.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.11.2 +importlib-metadata==4.11.3 # via keyring jeepney==0.7.1 # via @@ -68,9 +68,9 @@ jinja2-time==0.2.0 # via cookiecutter keyring==23.5.0 # via flytekit -markupsafe==2.1.0 +markupsafe==2.1.1 # via jinja2 -marshmallow==3.14.1 +marshmallow==3.15.0 # via # dataclasses-json # marshmallow-enum @@ -83,10 +83,12 @@ mypy-extensions==0.4.3 # via typing-inspect natsort==8.1.0 # via flytekit -numpy==1.22.2 +numpy==1.22.3 # via # pandas # pyarrow +packaging==21.3 + # via marshmallow pandas==1.4.1 # via flytekit poyo==0.5.0 @@ -105,6 +107,8 @@ pyarrow==6.0.1 # via flytekit pycparser==2.21 # via cffi +pyparsing==3.0.7 + # via packaging python-dateutil==2.8.2 # via # arrow @@ -121,14 +125,14 @@ pytz==2021.3 # via # flytekit # pandas -regex==2022.3.2 +regex==2022.3.15 # via docker-image-py requests==2.27.1 # via # cookiecutter # flytekit # responses -responses==0.18.0 +responses==0.19.0 # via flytekit retry==0.9.2 # via flytekit @@ -151,14 +155,14 @@ typing-extensions==4.1.1 # typing-inspect typing-inspect==0.7.1 # via dataclasses-json -urllib3==1.26.8 +urllib3==1.26.9 # via # flytekit # requests # responses wheel==0.37.1 # via flytekit -wrapt==1.13.3 +wrapt==1.14.0 # via # deprecated # flytekit diff --git a/plugins/flytekit-aws-batch/requirements.txt b/plugins/flytekit-aws-batch/requirements.txt index 80a3fcf51f..db5606322f 100644 --- a/plugins/flytekit-aws-batch/requirements.txt +++ b/plugins/flytekit-aws-batch/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.10 +# This file is autogenerated by pip-compile with python 3.8 # To update, run: # # pip-compile requirements.in @@ -30,7 +30,7 @@ cookiecutter==1.7.3 # via flytekit croniter==1.3.4 # via flytekit -cryptography==36.0.1 +cryptography==36.0.2 # via secretstorage dataclasses-json==0.5.6 # via flytekit @@ -44,17 +44,17 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.23.0 +flyteidl==0.24.0 # via flytekit -flytekit==0.30.3 +flytekit==0.31.0 # via flytekitplugins-awsbatch -googleapis-common-protos==1.55.0 +googleapis-common-protos==1.56.0 # via flyteidl grpcio==1.44.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.11.2 +importlib-metadata==4.11.3 # via keyring jeepney==0.7.1 # via @@ -68,9 +68,9 @@ jinja2-time==0.2.0 # via cookiecutter keyring==23.5.0 # via flytekit -markupsafe==2.1.0 +markupsafe==2.1.1 # via jinja2 -marshmallow==3.14.1 +marshmallow==3.15.0 # via # dataclasses-json # marshmallow-enum @@ -83,10 +83,12 @@ mypy-extensions==0.4.3 # via typing-inspect natsort==8.1.0 # via flytekit -numpy==1.22.2 +numpy==1.22.3 # via # pandas # pyarrow +packaging==21.3 + # via marshmallow pandas==1.4.1 # via flytekit poyo==0.5.0 @@ -105,6 +107,8 @@ pyarrow==6.0.1 # via flytekit pycparser==2.21 # via cffi +pyparsing==3.0.7 + # via packaging python-dateutil==2.8.2 # via # arrow @@ -121,14 +125,14 @@ pytz==2021.3 # via # flytekit # pandas -regex==2022.3.2 +regex==2022.3.15 # via docker-image-py requests==2.27.1 # via # cookiecutter # flytekit # responses -responses==0.18.0 +responses==0.19.0 # via flytekit retry==0.9.2 # via flytekit @@ -151,14 +155,14 @@ typing-extensions==4.1.1 # typing-inspect typing-inspect==0.7.1 # via dataclasses-json -urllib3==1.26.8 +urllib3==1.26.9 # via # flytekit # requests # responses wheel==0.37.1 # via flytekit -wrapt==1.13.3 +wrapt==1.14.0 # via # deprecated # flytekit diff --git a/plugins/flytekit-aws-sagemaker/requirements.txt b/plugins/flytekit-aws-sagemaker/requirements.txt index f4b02da229..54b04b5a13 100644 --- a/plugins/flytekit-aws-sagemaker/requirements.txt +++ b/plugins/flytekit-aws-sagemaker/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.10 +# This file is autogenerated by pip-compile with python 3.8 # To update, run: # # pip-compile requirements.in @@ -12,9 +12,9 @@ bcrypt==3.2.0 # via paramiko binaryornot==0.4.4 # via cookiecutter -boto3==1.21.10 +boto3==1.21.21 # via sagemaker-training -botocore==1.24.10 +botocore==1.24.21 # via # boto3 # s3transfer @@ -41,7 +41,7 @@ cookiecutter==1.7.3 # via flytekit croniter==1.3.4 # via flytekit -cryptography==36.0.1 +cryptography==36.0.2 # via # paramiko # secretstorage @@ -57,13 +57,13 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.23.0 +flyteidl==0.24.0 # via flytekit -flytekit==0.30.3 +flytekit==0.31.0 # via flytekitplugins-awssagemaker gevent==21.12.0 # via sagemaker-training -googleapis-common-protos==1.55.0 +googleapis-common-protos==1.56.0 # via flyteidl greenlet==1.1.2 # via gevent @@ -71,9 +71,9 @@ grpcio==1.44.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.11.2 +importlib-metadata==4.11.3 # via keyring -inotify_simple==1.2.1 +inotify-simple==1.2.1 # via sagemaker-training jeepney==0.7.1 # via @@ -85,15 +85,15 @@ jinja2==3.0.3 # jinja2-time jinja2-time==0.2.0 # via cookiecutter -jmespath==0.10.0 +jmespath==1.0.0 # via # boto3 # botocore keyring==23.5.0 # via flytekit -markupsafe==2.1.0 +markupsafe==2.1.1 # via jinja2 -marshmallow==3.14.1 +marshmallow==3.15.0 # via # dataclasses-json # marshmallow-enum @@ -106,15 +106,17 @@ mypy-extensions==0.4.3 # via typing-inspect natsort==8.1.0 # via flytekit -numpy==1.22.2 +numpy==1.22.3 # via # pandas # pyarrow # sagemaker-training # scipy +packaging==21.3 + # via marshmallow pandas==1.4.1 # via flytekit -paramiko==2.9.2 +paramiko==2.10.2 # via sagemaker-training poyo==0.5.0 # via cookiecutter @@ -137,6 +139,8 @@ pycparser==2.21 # via cffi pynacl==1.5.0 # via paramiko +pyparsing==3.0.7 + # via packaging python-dateutil==2.8.2 # via # arrow @@ -154,14 +158,14 @@ pytz==2021.3 # via # flytekit # pandas -regex==2022.3.2 +regex==2022.3.15 # via docker-image-py requests==2.27.1 # via # cookiecutter # flytekit # responses -responses==0.18.0 +responses==0.19.0 # via flytekit retry==0.9.2 # via flytekit @@ -180,6 +184,7 @@ six==1.16.0 # bcrypt # cookiecutter # grpcio + # paramiko # python-dateutil # retrying # sagemaker-training @@ -195,7 +200,7 @@ typing-extensions==4.1.1 # typing-inspect typing-inspect==0.7.1 # via dataclasses-json -urllib3==1.26.8 +urllib3==1.26.9 # via # botocore # flytekit @@ -205,15 +210,15 @@ werkzeug==2.0.3 # via sagemaker-training wheel==0.37.1 # via flytekit -wrapt==1.13.3 +wrapt==1.14.0 # via # deprecated # flytekit zipp==3.7.0 # via importlib-metadata -zope.event==4.5.0 +zope-event==4.5.0 # via gevent -zope.interface==5.4.0 +zope-interface==5.4.0 # via gevent # The following packages are considered to be unsafe in a requirements file: diff --git a/plugins/flytekit-bigquery/requirements.txt b/plugins/flytekit-bigquery/requirements.txt index aef0a38759..38576616bd 100644 --- a/plugins/flytekit-bigquery/requirements.txt +++ b/plugins/flytekit-bigquery/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.10 +# This file is autogenerated by pip-compile with python 3.8 # To update, run: # # pip-compile requirements.in @@ -32,7 +32,7 @@ cookiecutter==1.7.3 # via flytekit croniter==1.3.4 # via flytekit -cryptography==36.0.1 +cryptography==36.0.2 # via secretstorage dataclasses-json==0.5.6 # via flytekit @@ -46,27 +46,27 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.23.0 +flyteidl==0.24.0 # via flytekit -flytekit==0.30.3 +flytekit==0.31.0 # via flytekitplugins-bigquery -google-api-core[grpc]==2.5.0 +google-api-core[grpc]==2.7.1 # via # google-cloud-bigquery # google-cloud-core -google-auth==2.6.0 +google-auth==2.6.2 # via # google-api-core # google-cloud-core -google-cloud-bigquery==2.34.0 +google-cloud-bigquery==2.34.2 # via flytekitplugins-bigquery -google-cloud-core==2.2.2 +google-cloud-core==2.2.3 # via google-cloud-bigquery google-crc32c==1.3.0 # via google-resumable-media -google-resumable-media==2.3.0 +google-resumable-media==2.3.2 # via google-cloud-bigquery -googleapis-common-protos==1.55.0 +googleapis-common-protos==1.56.0 # via # flyteidl # google-api-core @@ -81,7 +81,7 @@ grpcio-status==1.44.0 # via google-api-core idna==3.3 # via requests -importlib-metadata==4.11.2 +importlib-metadata==4.11.3 # via keyring jeepney==0.7.1 # via @@ -95,9 +95,9 @@ jinja2-time==0.2.0 # via cookiecutter keyring==23.5.0 # via flytekit -markupsafe==2.1.0 +markupsafe==2.1.1 # via jinja2 -marshmallow==3.14.1 +marshmallow==3.15.0 # via # dataclasses-json # marshmallow-enum @@ -110,12 +110,14 @@ mypy-extensions==0.4.3 # via typing-inspect natsort==8.1.0 # via flytekit -numpy==1.22.2 +numpy==1.22.3 # via # pandas # pyarrow packaging==21.3 - # via google-cloud-bigquery + # via + # google-cloud-bigquery + # marshmallow pandas==1.4.1 # via flytekit poyo==0.5.0 @@ -165,7 +167,7 @@ pytz==2021.3 # via # flytekit # pandas -regex==2022.3.2 +regex==2022.3.15 # via docker-image-py requests==2.27.1 # via @@ -174,7 +176,7 @@ requests==2.27.1 # google-api-core # google-cloud-bigquery # responses -responses==0.18.0 +responses==0.19.0 # via flytekit retry==0.9.2 # via flytekit @@ -200,14 +202,14 @@ typing-extensions==4.1.1 # typing-inspect typing-inspect==0.7.1 # via dataclasses-json -urllib3==1.26.8 +urllib3==1.26.9 # via # flytekit # requests # responses wheel==0.37.1 # via flytekit -wrapt==1.13.3 +wrapt==1.14.0 # via # deprecated # flytekit diff --git a/plugins/flytekit-data-fsspec/requirements.txt b/plugins/flytekit-data-fsspec/requirements.txt index b3f506d572..f66ca64fba 100644 --- a/plugins/flytekit-data-fsspec/requirements.txt +++ b/plugins/flytekit-data-fsspec/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.10 +# This file is autogenerated by pip-compile with python 3.8 # To update, run: # # pip-compile requirements.in @@ -10,7 +10,7 @@ arrow==1.2.2 # via jinja2-time binaryornot==0.4.4 # via cookiecutter -botocore==1.24.10 +botocore==1.24.21 # via flytekitplugins-data-fsspec certifi==2021.10.8 # via requests @@ -32,7 +32,7 @@ cookiecutter==1.7.3 # via flytekit croniter==1.3.4 # via flytekit -cryptography==36.0.1 +cryptography==36.0.2 # via secretstorage dataclasses-json==0.5.6 # via flytekit @@ -46,19 +46,19 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.23.0 +flyteidl==0.24.0 # via flytekit -flytekit==0.30.3 +flytekit==0.31.0 # via flytekitplugins-data-fsspec fsspec==2022.2.0 # via flytekitplugins-data-fsspec -googleapis-common-protos==1.55.0 +googleapis-common-protos==1.56.0 # via flyteidl grpcio==1.44.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.11.2 +importlib-metadata==4.11.3 # via keyring jeepney==0.7.1 # via @@ -70,13 +70,13 @@ jinja2==3.0.3 # jinja2-time jinja2-time==0.2.0 # via cookiecutter -jmespath==0.10.0 +jmespath==1.0.0 # via botocore keyring==23.5.0 # via flytekit -markupsafe==2.1.0 +markupsafe==2.1.1 # via jinja2 -marshmallow==3.14.1 +marshmallow==3.15.0 # via # dataclasses-json # marshmallow-enum @@ -89,10 +89,12 @@ mypy-extensions==0.4.3 # via typing-inspect natsort==8.1.0 # via flytekit -numpy==1.22.2 +numpy==1.22.3 # via # pandas # pyarrow +packaging==21.3 + # via marshmallow pandas==1.4.1 # via flytekit poyo==0.5.0 @@ -111,6 +113,8 @@ pyarrow==6.0.1 # via flytekit pycparser==2.21 # via cffi +pyparsing==3.0.7 + # via packaging python-dateutil==2.8.2 # via # arrow @@ -128,14 +132,14 @@ pytz==2021.3 # via # flytekit # pandas -regex==2022.3.2 +regex==2022.3.15 # via docker-image-py requests==2.27.1 # via # cookiecutter # flytekit # responses -responses==0.18.0 +responses==0.19.0 # via flytekit retry==0.9.2 # via flytekit @@ -158,7 +162,7 @@ typing-extensions==4.1.1 # typing-inspect typing-inspect==0.7.1 # via dataclasses-json -urllib3==1.26.8 +urllib3==1.26.9 # via # botocore # flytekit @@ -166,7 +170,7 @@ urllib3==1.26.8 # responses wheel==0.37.1 # via flytekit -wrapt==1.13.3 +wrapt==1.14.0 # via # deprecated # flytekit diff --git a/plugins/flytekit-dolt/requirements.txt b/plugins/flytekit-dolt/requirements.txt index fc9dc9f8b7..bb61c3be8a 100644 --- a/plugins/flytekit-dolt/requirements.txt +++ b/plugins/flytekit-dolt/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.10 +# This file is autogenerated by pip-compile with python 3.8 # To update, run: # # pip-compile requirements.in @@ -30,7 +30,7 @@ cookiecutter==1.7.3 # via flytekit croniter==1.3.4 # via flytekit -cryptography==36.0.1 +cryptography==36.0.2 # via secretstorage dataclasses-json==0.5.6 # via @@ -50,17 +50,17 @@ dolt-integrations==0.1.5 # via flytekitplugins-dolt doltcli==0.1.17 # via dolt-integrations -flyteidl==0.23.0 +flyteidl==0.24.0 # via flytekit -flytekit==0.30.3 +flytekit==0.31.0 # via flytekitplugins-dolt -googleapis-common-protos==1.55.0 +googleapis-common-protos==1.56.0 # via flyteidl grpcio==1.44.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.11.2 +importlib-metadata==4.11.3 # via keyring jeepney==0.7.1 # via @@ -74,9 +74,9 @@ jinja2-time==0.2.0 # via cookiecutter keyring==23.5.0 # via flytekit -markupsafe==2.1.0 +markupsafe==2.1.1 # via jinja2 -marshmallow==3.14.1 +marshmallow==3.15.0 # via # dataclasses-json # marshmallow-enum @@ -89,10 +89,12 @@ mypy-extensions==0.4.3 # via typing-inspect natsort==8.1.0 # via flytekit -numpy==1.22.2 +numpy==1.22.3 # via # pandas # pyarrow +packaging==21.3 + # via marshmallow pandas==1.4.1 # via # dolt-integrations @@ -113,6 +115,8 @@ pyarrow==6.0.1 # via flytekit pycparser==2.21 # via cffi +pyparsing==3.0.7 + # via packaging python-dateutil==2.8.2 # via # arrow @@ -129,14 +133,14 @@ pytz==2021.3 # via # flytekit # pandas -regex==2022.3.2 +regex==2022.3.15 # via docker-image-py requests==2.27.1 # via # cookiecutter # flytekit # responses -responses==0.18.0 +responses==0.19.0 # via flytekit retry==0.9.2 # via flytekit @@ -159,14 +163,14 @@ typing-extensions==4.1.1 # typing-inspect typing-inspect==0.7.1 # via dataclasses-json -urllib3==1.26.8 +urllib3==1.26.9 # via # flytekit # requests # responses wheel==0.37.1 # via flytekit -wrapt==1.13.3 +wrapt==1.14.0 # via # deprecated # flytekit diff --git a/plugins/flytekit-greatexpectations/requirements.txt b/plugins/flytekit-greatexpectations/requirements.txt index 31b3be4b48..5c059de55f 100644 --- a/plugins/flytekit-greatexpectations/requirements.txt +++ b/plugins/flytekit-greatexpectations/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.10 +# This file is autogenerated by pip-compile with python 3.8 # To update, run: # # pip-compile requirements.in @@ -20,6 +20,12 @@ attrs==21.4.0 # via jsonschema backcall==0.2.0 # via ipython +backports-zoneinfo==0.2.1 + # via + # pytz-deprecation-shim + # tzlocal +beautifulsoup4==4.10.0 + # via nbconvert binaryornot==0.4.4 # via cookiecutter bleach==4.1.0 @@ -47,7 +53,7 @@ cookiecutter==1.7.3 # via flytekit croniter==1.3.4 # via flytekit -cryptography==36.0.1 +cryptography==36.0.2 # via secretstorage dataclasses-json==0.5.6 # via flytekit @@ -74,11 +80,11 @@ entrypoints==0.4 # nbconvert executing==0.8.3 # via stack-data -flyteidl==0.23.0 +flyteidl==0.24.0 # via flytekit -flytekit==0.30.3 +flytekit==0.31.0 # via flytekitplugins-great-expectations -googleapis-common-protos==1.55.0 +googleapis-common-protos==1.56.0 # via flyteidl great-expectations==0.14.5 # via flytekitplugins-great-expectations @@ -88,24 +94,25 @@ grpcio==1.44.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.11.2 +importlib-metadata==4.11.3 # via # great-expectations # keyring -ipykernel==6.9.1 +importlib-resources==5.4.0 + # via jsonschema +ipykernel==6.9.2 # via # ipywidgets # notebook -ipython==8.1.0 +ipython==8.1.1 # via # ipykernel # ipywidgets ipython-genutils==0.2.0 # via # ipywidgets - # nbformat # notebook -ipywidgets==7.6.5 +ipywidgets==7.7.0 # via great-expectations jedi==0.18.1 # via ipython @@ -145,13 +152,13 @@ jupyter-core==4.9.2 # notebook jupyterlab-pygments==0.1.2 # via nbconvert -jupyterlab-widgets==1.0.2 +jupyterlab-widgets==1.1.0 # via ipywidgets keyring==23.5.0 # via flytekit -markupsafe==2.1.0 +markupsafe==2.1.1 # via jinja2 -marshmallow==3.14.1 +marshmallow==3.15.0 # via # dataclasses-json # marshmallow-enum @@ -172,11 +179,11 @@ mypy-extensions==0.4.3 # via typing-inspect natsort==8.1.0 # via flytekit -nbclient==0.5.11 +nbclient==0.5.13 # via nbconvert -nbconvert==6.4.2 +nbconvert==6.4.4 # via notebook -nbformat==5.1.3 +nbformat==5.2.0 # via # ipywidgets # nbclient @@ -188,9 +195,9 @@ nest-asyncio==1.5.4 # jupyter-client # nbclient # notebook -notebook==6.4.8 +notebook==6.4.10 # via widgetsnbextension -numpy==1.22.2 +numpy==1.22.3 # via # altair # great-expectations @@ -198,7 +205,9 @@ numpy==1.22.2 # pyarrow # scipy packaging==21.3 - # via bleach + # via + # bleach + # marshmallow pandas==1.4.1 # via # altair @@ -226,6 +235,8 @@ protobuf==3.19.4 # protoc-gen-swagger protoc-gen-swagger==0.1.0 # via flyteidl +psutil==5.9.0 + # via ipykernel ptyprocess==0.7.0 # via # pexpect @@ -274,7 +285,7 @@ pyzmq==22.3.0 # via # jupyter-client # notebook -regex==2022.3.2 +regex==2022.3.15 # via docker-image-py requests==2.27.1 # via @@ -282,12 +293,14 @@ requests==2.27.1 # flytekit # great-expectations # responses -responses==0.18.0 +responses==0.19.0 # via flytekit retry==0.9.2 # via flytekit -ruamel.yaml==0.17.17 +ruamel-yaml==0.17.17 # via great-expectations +ruamel-yaml-clib==0.2.6 + # via ruamel-yaml scipy==1.8.0 # via great-expectations secretstorage==3.3.1 @@ -303,7 +316,9 @@ six==1.16.0 # python-dateutil sortedcontainers==2.4.0 # via flytekit -sqlalchemy==1.4.31 +soupsieve==2.3.1 + # via beautifulsoup4 +sqlalchemy==1.4.32 # via # -r requirements.in # flytekitplugins-great-expectations @@ -313,7 +328,7 @@ statsd==3.3.0 # via flytekit termcolor==1.1.0 # via great-expectations -terminado==0.13.1 +terminado==0.13.3 # via notebook testpath==0.6.0 # via nbconvert @@ -352,7 +367,7 @@ tzdata==2021.5 # via pytz-deprecation-shim tzlocal==4.1 # via great-expectations -urllib3==1.26.8 +urllib3==1.26.9 # via # flytekit # requests @@ -363,14 +378,16 @@ webencodings==0.5.1 # via bleach wheel==0.37.1 # via flytekit -widgetsnbextension==3.5.2 +widgetsnbextension==3.6.0 # via ipywidgets -wrapt==1.13.3 +wrapt==1.14.0 # via # deprecated # flytekit zipp==3.7.0 - # via importlib-metadata + # via + # importlib-metadata + # importlib-resources # The following packages are considered to be unsafe in a requirements file: # setuptools diff --git a/plugins/flytekit-hive/requirements.txt b/plugins/flytekit-hive/requirements.txt index 69ed1e5f3e..e52eb728d1 100644 --- a/plugins/flytekit-hive/requirements.txt +++ b/plugins/flytekit-hive/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.10 +# This file is autogenerated by pip-compile with python 3.8 # To update, run: # # pip-compile requirements.in @@ -30,7 +30,7 @@ cookiecutter==1.7.3 # via flytekit croniter==1.3.4 # via flytekit -cryptography==36.0.1 +cryptography==36.0.2 # via secretstorage dataclasses-json==0.5.6 # via flytekit @@ -44,17 +44,17 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.23.0 +flyteidl==0.24.0 # via flytekit -flytekit==0.30.3 +flytekit==0.31.0 # via flytekitplugins-hive -googleapis-common-protos==1.55.0 +googleapis-common-protos==1.56.0 # via flyteidl grpcio==1.44.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.11.2 +importlib-metadata==4.11.3 # via keyring jeepney==0.7.1 # via @@ -68,9 +68,9 @@ jinja2-time==0.2.0 # via cookiecutter keyring==23.5.0 # via flytekit -markupsafe==2.1.0 +markupsafe==2.1.1 # via jinja2 -marshmallow==3.14.1 +marshmallow==3.15.0 # via # dataclasses-json # marshmallow-enum @@ -83,10 +83,12 @@ mypy-extensions==0.4.3 # via typing-inspect natsort==8.1.0 # via flytekit -numpy==1.22.2 +numpy==1.22.3 # via # pandas # pyarrow +packaging==21.3 + # via marshmallow pandas==1.4.1 # via flytekit poyo==0.5.0 @@ -105,6 +107,8 @@ pyarrow==6.0.1 # via flytekit pycparser==2.21 # via cffi +pyparsing==3.0.7 + # via packaging python-dateutil==2.8.2 # via # arrow @@ -121,14 +125,14 @@ pytz==2021.3 # via # flytekit # pandas -regex==2022.3.2 +regex==2022.3.15 # via docker-image-py requests==2.27.1 # via # cookiecutter # flytekit # responses -responses==0.18.0 +responses==0.19.0 # via flytekit retry==0.9.2 # via flytekit @@ -151,14 +155,14 @@ typing-extensions==4.1.1 # typing-inspect typing-inspect==0.7.1 # via dataclasses-json -urllib3==1.26.8 +urllib3==1.26.9 # via # flytekit # requests # responses wheel==0.37.1 # via flytekit -wrapt==1.13.3 +wrapt==1.14.0 # via # deprecated # flytekit diff --git a/plugins/flytekit-k8s-pod/requirements.txt b/plugins/flytekit-k8s-pod/requirements.txt index df398913af..eb236d223f 100644 --- a/plugins/flytekit-k8s-pod/requirements.txt +++ b/plugins/flytekit-k8s-pod/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.10 +# This file is autogenerated by pip-compile with python 3.8 # To update, run: # # pip-compile requirements.in @@ -34,7 +34,7 @@ cookiecutter==1.7.3 # via flytekit croniter==1.3.4 # via flytekit -cryptography==36.0.1 +cryptography==36.0.2 # via secretstorage dataclasses-json==0.5.6 # via flytekit @@ -48,19 +48,19 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.23.0 +flyteidl==0.24.0 # via flytekit -flytekit==0.30.3 +flytekit==0.31.0 # via flytekitplugins-pod -google-auth==2.6.0 +google-auth==2.6.2 # via kubernetes -googleapis-common-protos==1.55.0 +googleapis-common-protos==1.56.0 # via flyteidl grpcio==1.44.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.11.2 +importlib-metadata==4.11.3 # via keyring jeepney==0.7.1 # via @@ -76,9 +76,9 @@ keyring==23.5.0 # via flytekit kubernetes==23.3.0 # via flytekitplugins-pod -markupsafe==2.1.0 +markupsafe==2.1.1 # via jinja2 -marshmallow==3.14.1 +marshmallow==3.15.0 # via # dataclasses-json # marshmallow-enum @@ -91,12 +91,14 @@ mypy-extensions==0.4.3 # via typing-inspect natsort==8.1.0 # via flytekit -numpy==1.22.2 +numpy==1.22.3 # via # pandas # pyarrow oauthlib==3.2.0 # via requests-oauthlib +packaging==21.3 + # via marshmallow pandas==1.4.1 # via flytekit poyo==0.5.0 @@ -121,6 +123,8 @@ pyasn1-modules==0.2.8 # via google-auth pycparser==2.21 # via cffi +pyparsing==3.0.7 + # via packaging python-dateutil==2.8.2 # via # arrow @@ -140,7 +144,7 @@ pytz==2021.3 # pandas pyyaml==6.0 # via kubernetes -regex==2022.3.2 +regex==2022.3.15 # via docker-image-py requests==2.27.1 # via @@ -151,7 +155,7 @@ requests==2.27.1 # responses requests-oauthlib==1.3.1 # via kubernetes -responses==0.18.0 +responses==0.19.0 # via flytekit retry==0.9.2 # via flytekit @@ -178,7 +182,7 @@ typing-extensions==4.1.1 # typing-inspect typing-inspect==0.7.1 # via dataclasses-json -urllib3==1.26.8 +urllib3==1.26.9 # via # flytekit # kubernetes @@ -188,7 +192,7 @@ websocket-client==1.3.1 # via kubernetes wheel==0.37.1 # via flytekit -wrapt==1.13.3 +wrapt==1.14.0 # via # deprecated # flytekit diff --git a/plugins/flytekit-kf-mpi/requirements.txt b/plugins/flytekit-kf-mpi/requirements.txt index 82a2317f2f..6af7ac0ab6 100644 --- a/plugins/flytekit-kf-mpi/requirements.txt +++ b/plugins/flytekit-kf-mpi/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.10 +# This file is autogenerated by pip-compile with python 3.8 # To update, run: # # pip-compile requirements.in @@ -30,7 +30,7 @@ cookiecutter==1.7.3 # via flytekit croniter==1.3.4 # via flytekit -cryptography==36.0.1 +cryptography==36.0.2 # via secretstorage dataclasses-json==0.5.6 # via flytekit @@ -44,19 +44,19 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.23.0 +flyteidl==0.24.0 # via # flytekit # flytekitplugins-kfmpi -flytekit==0.30.3 +flytekit==0.31.0 # via flytekitplugins-kfmpi -googleapis-common-protos==1.55.0 +googleapis-common-protos==1.56.0 # via flyteidl grpcio==1.44.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.11.2 +importlib-metadata==4.11.3 # via keyring jeepney==0.7.1 # via @@ -70,9 +70,9 @@ jinja2-time==0.2.0 # via cookiecutter keyring==23.5.0 # via flytekit -markupsafe==2.1.0 +markupsafe==2.1.1 # via jinja2 -marshmallow==3.14.1 +marshmallow==3.15.0 # via # dataclasses-json # marshmallow-enum @@ -85,10 +85,12 @@ mypy-extensions==0.4.3 # via typing-inspect natsort==8.1.0 # via flytekit -numpy==1.22.2 +numpy==1.22.3 # via # pandas # pyarrow +packaging==21.3 + # via marshmallow pandas==1.4.1 # via flytekit poyo==0.5.0 @@ -107,6 +109,8 @@ pyarrow==6.0.1 # via flytekit pycparser==2.21 # via cffi +pyparsing==3.0.7 + # via packaging python-dateutil==2.8.2 # via # arrow @@ -123,14 +127,14 @@ pytz==2021.3 # via # flytekit # pandas -regex==2022.3.2 +regex==2022.3.15 # via docker-image-py requests==2.27.1 # via # cookiecutter # flytekit # responses -responses==0.18.0 +responses==0.19.0 # via flytekit retry==0.9.2 # via flytekit @@ -153,14 +157,14 @@ typing-extensions==4.1.1 # typing-inspect typing-inspect==0.7.1 # via dataclasses-json -urllib3==1.26.8 +urllib3==1.26.9 # via # flytekit # requests # responses wheel==0.37.1 # via flytekit -wrapt==1.13.3 +wrapt==1.14.0 # via # deprecated # flytekit diff --git a/plugins/flytekit-kf-pytorch/requirements.txt b/plugins/flytekit-kf-pytorch/requirements.txt index 0a53a62333..5354e80db2 100644 --- a/plugins/flytekit-kf-pytorch/requirements.txt +++ b/plugins/flytekit-kf-pytorch/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.10 +# This file is autogenerated by pip-compile with python 3.8 # To update, run: # # pip-compile requirements.in @@ -30,7 +30,7 @@ cookiecutter==1.7.3 # via flytekit croniter==1.3.4 # via flytekit -cryptography==36.0.1 +cryptography==36.0.2 # via secretstorage dataclasses-json==0.5.6 # via flytekit @@ -44,17 +44,17 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.23.0 +flyteidl==0.24.0 # via flytekit -flytekit==0.30.3 +flytekit==0.31.0 # via flytekitplugins-kfpytorch -googleapis-common-protos==1.55.0 +googleapis-common-protos==1.56.0 # via flyteidl grpcio==1.44.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.11.2 +importlib-metadata==4.11.3 # via keyring jeepney==0.7.1 # via @@ -68,9 +68,9 @@ jinja2-time==0.2.0 # via cookiecutter keyring==23.5.0 # via flytekit -markupsafe==2.1.0 +markupsafe==2.1.1 # via jinja2 -marshmallow==3.14.1 +marshmallow==3.15.0 # via # dataclasses-json # marshmallow-enum @@ -83,10 +83,12 @@ mypy-extensions==0.4.3 # via typing-inspect natsort==8.1.0 # via flytekit -numpy==1.22.2 +numpy==1.22.3 # via # pandas # pyarrow +packaging==21.3 + # via marshmallow pandas==1.4.1 # via flytekit poyo==0.5.0 @@ -105,6 +107,8 @@ pyarrow==6.0.1 # via flytekit pycparser==2.21 # via cffi +pyparsing==3.0.7 + # via packaging python-dateutil==2.8.2 # via # arrow @@ -121,14 +125,14 @@ pytz==2021.3 # via # flytekit # pandas -regex==2022.3.2 +regex==2022.3.15 # via docker-image-py requests==2.27.1 # via # cookiecutter # flytekit # responses -responses==0.18.0 +responses==0.19.0 # via flytekit retry==0.9.2 # via flytekit @@ -151,14 +155,14 @@ typing-extensions==4.1.1 # typing-inspect typing-inspect==0.7.1 # via dataclasses-json -urllib3==1.26.8 +urllib3==1.26.9 # via # flytekit # requests # responses wheel==0.37.1 # via flytekit -wrapt==1.13.3 +wrapt==1.14.0 # via # deprecated # flytekit diff --git a/plugins/flytekit-kf-tensorflow/requirements.txt b/plugins/flytekit-kf-tensorflow/requirements.txt index b782c39630..749255fecf 100644 --- a/plugins/flytekit-kf-tensorflow/requirements.txt +++ b/plugins/flytekit-kf-tensorflow/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.10 +# This file is autogenerated by pip-compile with python 3.8 # To update, run: # # pip-compile requirements.in @@ -30,7 +30,7 @@ cookiecutter==1.7.3 # via flytekit croniter==1.3.4 # via flytekit -cryptography==36.0.1 +cryptography==36.0.2 # via secretstorage dataclasses-json==0.5.6 # via flytekit @@ -44,17 +44,17 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.23.0 +flyteidl==0.24.0 # via flytekit -flytekit==0.30.3 +flytekit==0.31.0 # via flytekitplugins-kftensorflow -googleapis-common-protos==1.55.0 +googleapis-common-protos==1.56.0 # via flyteidl grpcio==1.44.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.11.2 +importlib-metadata==4.11.3 # via keyring jeepney==0.7.1 # via @@ -68,9 +68,9 @@ jinja2-time==0.2.0 # via cookiecutter keyring==23.5.0 # via flytekit -markupsafe==2.1.0 +markupsafe==2.1.1 # via jinja2 -marshmallow==3.14.1 +marshmallow==3.15.0 # via # dataclasses-json # marshmallow-enum @@ -83,10 +83,12 @@ mypy-extensions==0.4.3 # via typing-inspect natsort==8.1.0 # via flytekit -numpy==1.22.2 +numpy==1.22.3 # via # pandas # pyarrow +packaging==21.3 + # via marshmallow pandas==1.4.1 # via flytekit poyo==0.5.0 @@ -105,6 +107,8 @@ pyarrow==6.0.1 # via flytekit pycparser==2.21 # via cffi +pyparsing==3.0.7 + # via packaging python-dateutil==2.8.2 # via # arrow @@ -121,14 +125,14 @@ pytz==2021.3 # via # flytekit # pandas -regex==2022.3.2 +regex==2022.3.15 # via docker-image-py requests==2.27.1 # via # cookiecutter # flytekit # responses -responses==0.18.0 +responses==0.19.0 # via flytekit retry==0.9.2 # via flytekit @@ -151,14 +155,14 @@ typing-extensions==4.1.1 # typing-inspect typing-inspect==0.7.1 # via dataclasses-json -urllib3==1.26.8 +urllib3==1.26.9 # via # flytekit # requests # responses wheel==0.37.1 # via flytekit -wrapt==1.13.3 +wrapt==1.14.0 # via # deprecated # flytekit diff --git a/plugins/flytekit-modin/requirements.txt b/plugins/flytekit-modin/requirements.txt index 4f3b01e5e4..fefbb4d944 100644 --- a/plugins/flytekit-modin/requirements.txt +++ b/plugins/flytekit-modin/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.9 +# This file is autogenerated by pip-compile with python 3.8 # To update, run: # # pip-compile requirements.in @@ -16,13 +16,15 @@ binaryornot==0.4.4 # via cookiecutter certifi==2021.10.8 # via requests +cffi==1.15.0 + # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.11 +charset-normalizer==2.0.12 # via requests checksumdir==1.2.0 # via flytekit -click==7.1.2 +click==8.0.4 # via # cookiecutter # flytekit @@ -31,9 +33,9 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.2.0 +croniter==1.3.4 # via flytekit -cryptography==36.0.1 +cryptography==36.0.2 # via secretstorage dataclasses-json==0.5.6 # via flytekit @@ -49,24 +51,32 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -filelock==3.4.2 +filelock==3.6.0 # via ray -flyteidl==0.22.0 +flyteidl==0.24.0 # via flytekit -flytekit==0.30.0 +flytekit==0.31.0 # via flytekitplugins-modin -fsspec==2022.1.0 +fsspec==2022.2.0 # via # flytekitplugins-modin # modin +googleapis-common-protos==1.56.0 + # via flyteidl grpcio==1.43.0 # via # flytekit # ray idna==3.3 # via requests -importlib-metadata==4.10.1 +importlib-metadata==4.11.3 # via keyring +importlib-resources==5.4.0 + # via jsonschema +jeepney==0.7.1 + # via + # keyring + # secretstorage jinja2==3.0.3 # via # cookiecutter @@ -77,9 +87,9 @@ jsonschema==4.4.0 # via ray keyring==23.5.0 # via flytekit -markupsafe==2.0.1 +markupsafe==2.1.1 # via jinja2 -marshmallow==3.14.1 +marshmallow==3.15.0 # via # dataclasses-json # marshmallow-enum @@ -88,7 +98,7 @@ marshmallow-enum==1.5.1 # via dataclasses-json marshmallow-jsonschema==0.13.0 # via flytekit -modin==0.13.1 +modin==0.13.2 # via flytekitplugins-modin msgpack==1.0.3 # via ray @@ -96,7 +106,7 @@ mypy-extensions==0.4.3 # via typing-inspect natsort==8.1.0 # via flytekit -numpy==1.22.2 +numpy==1.22.3 # via # modin # pandas @@ -104,6 +114,7 @@ numpy==1.22.2 # ray packaging==21.3 # via + # marshmallow # modin # redis pandas==1.4.0 @@ -116,7 +127,11 @@ protobuf==3.19.4 # via # flyteidl # flytekit + # googleapis-common-protos + # protoc-gen-swagger # ray +protoc-gen-swagger==0.1.0 + # via flyteidl py==1.11.0 # via retry pyarrow==6.0.1 @@ -135,7 +150,7 @@ python-dateutil==2.8.2 # pandas python-json-logger==2.0.2 # via flytekit -python-slugify==5.0.2 +python-slugify==6.1.1 # via cookiecutter pytimeparse==1.1.8 # via flytekit @@ -145,21 +160,23 @@ pytz==2021.3 # pandas pyyaml==6.0 # via ray -ray==1.10.0 +ray==1.11.0 # via flytekitplugins-modin -redis==4.1.2 +redis==4.1.4 # via ray -regex==2022.1.18 +regex==2022.3.15 # via docker-image-py requests==2.27.1 # via # cookiecutter # flytekit # responses -responses==0.18.0 +responses==0.19.0 # via flytekit retry==0.9.2 # via flytekit +secretstorage==3.3.1 + # via keyring six==1.16.0 # via # cookiecutter @@ -171,22 +188,24 @@ statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -typing-extensions==4.0.1 +typing-extensions==4.1.1 # via # flytekit # typing-inspect typing-inspect==0.7.1 # via dataclasses-json -urllib3==1.26.8 +urllib3==1.26.9 # via # flytekit # requests # responses wheel==0.37.1 # via flytekit -wrapt==1.13.3 +wrapt==1.14.0 # via # deprecated # flytekit zipp==3.7.0 - # via importlib-metadata + # via + # importlib-metadata + # importlib-resources diff --git a/plugins/flytekit-pandera/requirements.txt b/plugins/flytekit-pandera/requirements.txt index 7b72a0b4c0..a4c86dba63 100644 --- a/plugins/flytekit-pandera/requirements.txt +++ b/plugins/flytekit-pandera/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.10 +# This file is autogenerated by pip-compile with python 3.8 # To update, run: # # pip-compile requirements.in @@ -30,7 +30,7 @@ cookiecutter==1.7.3 # via flytekit croniter==1.3.4 # via flytekit -cryptography==36.0.1 +cryptography==36.0.2 # via secretstorage dataclasses-json==0.5.6 # via flytekit @@ -44,17 +44,17 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.23.0 +flyteidl==0.24.0 # via flytekit -flytekit==0.30.3 +flytekit==0.31.0 # via flytekitplugins-pandera -googleapis-common-protos==1.55.0 +googleapis-common-protos==1.56.0 # via flyteidl grpcio==1.44.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.11.2 +importlib-metadata==4.11.3 # via keyring jeepney==0.7.1 # via @@ -68,9 +68,9 @@ jinja2-time==0.2.0 # via cookiecutter keyring==23.5.0 # via flytekit -markupsafe==2.1.0 +markupsafe==2.1.1 # via jinja2 -marshmallow==3.14.1 +marshmallow==3.15.0 # via # dataclasses-json # marshmallow-enum @@ -83,13 +83,15 @@ mypy-extensions==0.4.3 # via typing-inspect natsort==8.1.0 # via flytekit -numpy==1.22.2 +numpy==1.22.3 # via # pandas # pandera # pyarrow packaging==21.3 - # via pandera + # via + # marshmallow + # pandera pandas==1.4.1 # via # flytekit @@ -134,14 +136,14 @@ pytz==2021.3 # via # flytekit # pandas -regex==2022.3.2 +regex==2022.3.15 # via docker-image-py requests==2.27.1 # via # cookiecutter # flytekit # responses -responses==0.18.0 +responses==0.19.0 # via flytekit retry==0.9.2 # via flytekit @@ -167,14 +169,14 @@ typing-inspect==0.7.1 # via # dataclasses-json # pandera -urllib3==1.26.8 +urllib3==1.26.9 # via # flytekit # requests # responses wheel==0.37.1 # via flytekit -wrapt==1.13.3 +wrapt==1.14.0 # via # deprecated # flytekit diff --git a/plugins/flytekit-papermill/requirements.txt b/plugins/flytekit-papermill/requirements.txt index 8c3e4570e9..9463e5cbd7 100644 --- a/plugins/flytekit-papermill/requirements.txt +++ b/plugins/flytekit-papermill/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.9 +# This file is autogenerated by pip-compile with python 3.8 # To update, run: # # pip-compile requirements.in @@ -16,10 +16,10 @@ attrs==21.4.0 # via jsonschema backcall==0.2.0 # via ipython +beautifulsoup4==4.10.0 + # via nbconvert binaryornot==0.4.4 # via cookiecutter -black==21.12b0 - # via ipython bleach==4.1.0 # via nbconvert certifi==2021.10.8 @@ -28,13 +28,12 @@ cffi==1.15.0 # via cryptography chardet==4.0.0 # via binaryornot -charset-normalizer==2.0.10 +charset-normalizer==2.0.12 # via requests checksumdir==1.2.0 # via flytekit -click==7.1.2 +click==8.0.4 # via - # black # cookiecutter # flytekit # papermill @@ -42,9 +41,9 @@ cloudpickle==2.0.0 # via flytekit cookiecutter==1.7.3 # via flytekit -croniter==1.2.0 +croniter==1.3.4 # via flytekit -cryptography==36.0.1 +cryptography==36.0.2 # via secretstorage dataclasses-json==0.5.6 # via flytekit @@ -64,29 +63,31 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -entrypoints==0.3 +entrypoints==0.4 # via # jupyter-client # nbconvert # papermill -executing==0.8.2 +executing==0.8.3 # via stack-data -flyteidl==0.22.0 +flyteidl==0.24.0 # via flytekit -flytekit==0.26.1 +flytekit==0.31.0 # via flytekitplugins-papermill -grpcio==1.43.0 +googleapis-common-protos==1.56.0 + # via flyteidl +grpcio==1.44.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.10.1 +importlib-metadata==4.11.3 # via keyring -ipykernel==6.7.0 +importlib-resources==5.4.0 + # via jsonschema +ipykernel==6.9.2 # via flytekitplugins-papermill -ipython==8.0.1 +ipython==8.1.1 # via ipykernel -ipython-genutils==0.2.0 - # via nbformat jedi==0.18.1 # via ipython jeepney==0.7.1 @@ -106,7 +107,7 @@ jupyter-client==7.1.2 # via # ipykernel # nbclient -jupyter-core==4.9.1 +jupyter-core==4.9.2 # via # jupyter-client # nbconvert @@ -115,9 +116,9 @@ jupyterlab-pygments==0.1.2 # via nbconvert keyring==23.5.0 # via flytekit -markupsafe==2.0.1 +markupsafe==2.1.1 # via jinja2 -marshmallow==3.14.1 +marshmallow==3.15.0 # via # dataclasses-json # marshmallow-enum @@ -133,18 +134,16 @@ matplotlib-inline==0.1.3 mistune==0.8.4 # via nbconvert mypy-extensions==0.4.3 - # via - # black - # typing-inspect -natsort==8.0.2 + # via typing-inspect +natsort==8.1.0 # via flytekit -nbclient==0.5.10 +nbclient==0.5.13 # via # nbconvert # papermill -nbconvert==6.4.1 +nbconvert==6.4.4 # via flytekitplugins-papermill -nbformat==5.1.3 +nbformat==5.2.0 # via # nbclient # nbconvert @@ -154,13 +153,15 @@ nest-asyncio==1.5.4 # ipykernel # jupyter-client # nbclient -numpy==1.22.1 +numpy==1.22.3 # via # pandas # pyarrow packaging==21.3 - # via bleach -pandas==1.4.0 + # via + # bleach + # marshmallow +pandas==1.4.1 # via flytekit pandocfilters==1.5.0 # via nbconvert @@ -168,22 +169,24 @@ papermill==2.3.4 # via flytekitplugins-papermill parso==0.8.3 # via jedi -pathspec==0.9.0 - # via black pexpect==4.8.0 # via ipython pickleshare==0.7.5 # via ipython -platformdirs==2.4.1 - # via black poyo==0.5.0 # via cookiecutter -prompt-toolkit==3.0.26 +prompt-toolkit==3.0.28 # via ipython protobuf==3.19.4 # via # flyteidl # flytekit + # googleapis-common-protos + # protoc-gen-swagger +protoc-gen-swagger==0.1.0 + # via flyteidl +psutil==5.9.0 + # via ipykernel ptyprocess==0.7.0 # via pexpect pure-eval==0.2.2 @@ -203,7 +206,7 @@ pyparsing==3.0.7 # via packaging pyrsistent==0.18.1 # via jsonschema -python-dateutil==2.8.1 +python-dateutil==2.8.2 # via # arrow # croniter @@ -212,7 +215,7 @@ python-dateutil==2.8.1 # pandas python-json-logger==2.0.2 # via flytekit -python-slugify==5.0.2 +python-slugify==6.1.1 # via cookiecutter pytimeparse==1.1.8 # via flytekit @@ -224,7 +227,7 @@ pyyaml==6.0 # via papermill pyzmq==22.3.0 # via jupyter-client -regex==2022.1.18 +regex==2022.3.15 # via docker-image-py requests==2.27.1 # via @@ -232,7 +235,7 @@ requests==2.27.1 # flytekit # papermill # responses -responses==0.17.0 +responses==0.19.0 # via flytekit retry==0.9.2 # via flytekit @@ -243,31 +246,29 @@ six==1.16.0 # asttokens # bleach # cookiecutter - # flytekit # grpcio # python-dateutil - # responses sortedcontainers==2.4.0 # via flytekit -stack-data==0.1.4 +soupsieve==2.3.1 + # via beautifulsoup4 +stack-data==0.2.0 # via ipython statsd==3.3.0 # via flytekit tenacity==8.0.1 # via papermill -testpath==0.5.0 +testpath==0.6.0 # via nbconvert text-unidecode==1.3 # via python-slugify textwrap3==0.9.2 # via ansiwrap -tomli==1.2.3 - # via black tornado==6.1 # via # ipykernel # jupyter-client -tqdm==4.62.3 +tqdm==4.63.0 # via papermill traitlets==5.1.1 # via @@ -279,13 +280,13 @@ traitlets==5.1.1 # nbclient # nbconvert # nbformat -typing-extensions==4.0.1 +typing-extensions==4.1.1 # via - # black + # flytekit # typing-inspect typing-inspect==0.7.1 # via dataclasses-json -urllib3==1.26.8 +urllib3==1.26.9 # via # flytekit # requests @@ -296,12 +297,14 @@ webencodings==0.5.1 # via bleach wheel==0.37.1 # via flytekit -wrapt==1.13.3 +wrapt==1.14.0 # via # deprecated # flytekit zipp==3.7.0 - # via importlib-metadata + # via + # importlib-metadata + # importlib-resources # The following packages are considered to be unsafe in a requirements file: # setuptools diff --git a/plugins/flytekit-snowflake/requirements.txt b/plugins/flytekit-snowflake/requirements.txt index 2d34ec966f..43c8f445e7 100644 --- a/plugins/flytekit-snowflake/requirements.txt +++ b/plugins/flytekit-snowflake/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.10 +# This file is autogenerated by pip-compile with python 3.8 # To update, run: # # pip-compile requirements.in @@ -30,7 +30,7 @@ cookiecutter==1.7.3 # via flytekit croniter==1.3.4 # via flytekit -cryptography==36.0.1 +cryptography==36.0.2 # via secretstorage dataclasses-json==0.5.6 # via flytekit @@ -44,17 +44,17 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.23.0 +flyteidl==0.24.0 # via flytekit -flytekit==0.30.3 +flytekit==0.31.0 # via flytekitplugins-snowflake -googleapis-common-protos==1.55.0 +googleapis-common-protos==1.56.0 # via flyteidl grpcio==1.44.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.11.2 +importlib-metadata==4.11.3 # via keyring jeepney==0.7.1 # via @@ -68,9 +68,9 @@ jinja2-time==0.2.0 # via cookiecutter keyring==23.5.0 # via flytekit -markupsafe==2.1.0 +markupsafe==2.1.1 # via jinja2 -marshmallow==3.14.1 +marshmallow==3.15.0 # via # dataclasses-json # marshmallow-enum @@ -83,10 +83,12 @@ mypy-extensions==0.4.3 # via typing-inspect natsort==8.1.0 # via flytekit -numpy==1.22.2 +numpy==1.22.3 # via # pandas # pyarrow +packaging==21.3 + # via marshmallow pandas==1.4.1 # via flytekit poyo==0.5.0 @@ -105,6 +107,8 @@ pyarrow==6.0.1 # via flytekit pycparser==2.21 # via cffi +pyparsing==3.0.7 + # via packaging python-dateutil==2.8.2 # via # arrow @@ -121,14 +125,14 @@ pytz==2021.3 # via # flytekit # pandas -regex==2022.3.2 +regex==2022.3.15 # via docker-image-py requests==2.27.1 # via # cookiecutter # flytekit # responses -responses==0.18.0 +responses==0.19.0 # via flytekit retry==0.9.2 # via flytekit @@ -151,14 +155,14 @@ typing-extensions==4.1.1 # typing-inspect typing-inspect==0.7.1 # via dataclasses-json -urllib3==1.26.8 +urllib3==1.26.9 # via # flytekit # requests # responses wheel==0.37.1 # via flytekit -wrapt==1.13.3 +wrapt==1.14.0 # via # deprecated # flytekit diff --git a/plugins/flytekit-spark/requirements.txt b/plugins/flytekit-spark/requirements.txt index 6e7c6e59f0..f53e809b77 100644 --- a/plugins/flytekit-spark/requirements.txt +++ b/plugins/flytekit-spark/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.10 +# This file is autogenerated by pip-compile with python 3.8 # To update, run: # # pip-compile requirements.in @@ -30,7 +30,7 @@ cookiecutter==1.7.3 # via flytekit croniter==1.3.4 # via flytekit -cryptography==36.0.1 +cryptography==36.0.2 # via secretstorage dataclasses-json==0.5.6 # via flytekit @@ -44,17 +44,17 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.23.0 +flyteidl==0.24.0 # via flytekit -flytekit==0.30.3 +flytekit==0.31.0 # via flytekitplugins-spark -googleapis-common-protos==1.55.0 +googleapis-common-protos==1.56.0 # via flyteidl grpcio==1.44.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.11.2 +importlib-metadata==4.11.3 # via keyring jeepney==0.7.1 # via @@ -68,9 +68,9 @@ jinja2-time==0.2.0 # via cookiecutter keyring==23.5.0 # via flytekit -markupsafe==2.1.0 +markupsafe==2.1.1 # via jinja2 -marshmallow==3.14.1 +marshmallow==3.15.0 # via # dataclasses-json # marshmallow-enum @@ -83,10 +83,12 @@ mypy-extensions==0.4.3 # via typing-inspect natsort==8.1.0 # via flytekit -numpy==1.22.2 +numpy==1.22.3 # via # pandas # pyarrow +packaging==21.3 + # via marshmallow pandas==1.4.1 # via flytekit poyo==0.5.0 @@ -107,6 +109,8 @@ pyarrow==6.0.1 # via flytekit pycparser==2.21 # via cffi +pyparsing==3.0.7 + # via packaging pyspark==3.2.1 # via flytekitplugins-spark python-dateutil==2.8.2 @@ -125,14 +129,14 @@ pytz==2021.3 # via # flytekit # pandas -regex==2022.3.2 +regex==2022.3.15 # via docker-image-py requests==2.27.1 # via # cookiecutter # flytekit # responses -responses==0.18.0 +responses==0.19.0 # via flytekit retry==0.9.2 # via flytekit @@ -155,14 +159,14 @@ typing-extensions==4.1.1 # typing-inspect typing-inspect==0.7.1 # via dataclasses-json -urllib3==1.26.8 +urllib3==1.26.9 # via # flytekit # requests # responses wheel==0.37.1 # via flytekit -wrapt==1.13.3 +wrapt==1.14.0 # via # deprecated # flytekit diff --git a/plugins/flytekit-sqlalchemy/requirements.txt b/plugins/flytekit-sqlalchemy/requirements.txt index e617920549..77485ea9cb 100644 --- a/plugins/flytekit-sqlalchemy/requirements.txt +++ b/plugins/flytekit-sqlalchemy/requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.10 +# This file is autogenerated by pip-compile with python 3.8 # To update, run: # # pip-compile requirements.in @@ -30,7 +30,7 @@ cookiecutter==1.7.3 # via flytekit croniter==1.3.4 # via flytekit -cryptography==36.0.1 +cryptography==36.0.2 # via secretstorage dataclasses-json==0.5.6 # via flytekit @@ -44,11 +44,11 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.23.0 +flyteidl==0.24.0 # via flytekit -flytekit==0.30.3 +flytekit==0.31.0 # via flytekitplugins-sqlalchemy -googleapis-common-protos==1.55.0 +googleapis-common-protos==1.56.0 # via flyteidl greenlet==1.1.2 # via sqlalchemy @@ -56,7 +56,7 @@ grpcio==1.44.0 # via flytekit idna==3.3 # via requests -importlib-metadata==4.11.2 +importlib-metadata==4.11.3 # via keyring jeepney==0.7.1 # via @@ -70,9 +70,9 @@ jinja2-time==0.2.0 # via cookiecutter keyring==23.5.0 # via flytekit -markupsafe==2.1.0 +markupsafe==2.1.1 # via jinja2 -marshmallow==3.14.1 +marshmallow==3.15.0 # via # dataclasses-json # marshmallow-enum @@ -85,10 +85,12 @@ mypy-extensions==0.4.3 # via typing-inspect natsort==8.1.0 # via flytekit -numpy==1.22.2 +numpy==1.22.3 # via # pandas # pyarrow +packaging==21.3 + # via marshmallow pandas==1.4.1 # via flytekit poyo==0.5.0 @@ -107,6 +109,8 @@ pyarrow==6.0.1 # via flytekit pycparser==2.21 # via cffi +pyparsing==3.0.7 + # via packaging python-dateutil==2.8.2 # via # arrow @@ -123,14 +127,14 @@ pytz==2021.3 # via # flytekit # pandas -regex==2022.3.2 +regex==2022.3.15 # via docker-image-py requests==2.27.1 # via # cookiecutter # flytekit # responses -responses==0.18.0 +responses==0.19.0 # via flytekit retry==0.9.2 # via flytekit @@ -143,7 +147,7 @@ six==1.16.0 # python-dateutil sortedcontainers==2.4.0 # via flytekit -sqlalchemy==1.4.31 +sqlalchemy==1.4.32 # via flytekitplugins-sqlalchemy statsd==3.3.0 # via flytekit @@ -155,14 +159,14 @@ typing-extensions==4.1.1 # typing-inspect typing-inspect==0.7.1 # via dataclasses-json -urllib3==1.26.8 +urllib3==1.26.9 # via # flytekit # requests # responses wheel==0.37.1 # via flytekit -wrapt==1.13.3 +wrapt==1.14.0 # via # deprecated # flytekit From e8610db328e69e2df2bad6a2efb6cbdd637830e4 Mon Sep 17 00:00:00 2001 From: Eduardo Apolinario Date: Thu, 17 Mar 2022 17:21:58 -0700 Subject: [PATCH 119/128] Handle nested Annotated Signed-off-by: Eduardo Apolinario --- flytekit/core/type_engine.py | 2 +- tests/flytekit/unit/core/test_type_engine.py | 14 ++++++++++++++ 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 5fca8cd654..87160b9e83 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -171,7 +171,7 @@ def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], exp return self._to_literal_transformer(python_val) def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[T]) -> T: - while get_origin(expected_python_type) is Annotated: + if get_origin(expected_python_type) is Annotated: expected_python_type = get_args(expected_python_type)[0] if expected_python_type != self._type: diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 639c549f91..4967f7cda5 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -1141,6 +1141,20 @@ def test_dict_to_literal_map_with_wrong_input_type(): TypeEngine.dict_to_literal_map(ctx, input, guessed_python_types) +def test_nested_annotated(): + """ + Test to show that nested Annotated types are flattened. + """ + pt = Annotated[Annotated[int, 'inner-annotation'], 'outer-annotation'] + lt = TypeEngine.to_literal_type(pt) + assert lt.simple == model_types.SimpleType.INTEGER + + ctx = FlyteContextManager.current_context() + lv = TypeEngine.to_literal(ctx, 42, pt, lt) + v = TypeEngine.to_python_value(ctx, lv, pt) + assert v == 42 + + def test_pass_annotated_to_downstream_tasks(): """ Test to confirm that the loaded dataframe is not affected and can be used in @dynamic. From b8500927eed7631eb4d6d43138f454038ab9e39b Mon Sep 17 00:00:00 2001 From: Eduardo Apolinario Date: Thu, 17 Mar 2022 17:27:36 -0700 Subject: [PATCH 120/128] Leave TODO re: strucutured dataset type castability Signed-off-by: Eduardo Apolinario --- flytekit/core/type_engine.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 87160b9e83..6e1dd9f043 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -19,6 +19,7 @@ from google.protobuf.struct_pb2 import Struct from marshmallow_enum import EnumField, LoadDumpOptions from marshmallow_jsonschema import JSONSchema +from typing_extensions import Annotated, get_args, get_origin from flytekit.core.annotation import FlyteAnnotation from flytekit.core.context_manager import FlyteContext @@ -45,11 +46,6 @@ ) from flytekit.models.types import LiteralType, SimpleType, StructuredDatasetType, TypeStructure, UnionType -try: - from typing import Annotated, get_args, get_origin -except ImportError: - from typing_extensions import Annotated, get_args, get_origin - T = typing.TypeVar("T") DEFINITIONS = "definitions" @@ -875,6 +871,8 @@ def _are_types_castable(upstream: LiteralType, downstream: LiteralType) -> bool: return _are_types_castable(upstream.map_value_type, downstream.map_value_type) + # TODO: Structured dataset type matching requires that downstream structured datasets + # are a strict sub-set of the upstream structured dataset. if upstream.structured_dataset_type is not None: if downstream.structured_dataset_type is None: return False From 6587c8fafdb73a86fb510754a946f73f5e562eba Mon Sep 17 00:00:00 2001 From: Eduardo Apolinario Date: Thu, 17 Mar 2022 17:30:46 -0700 Subject: [PATCH 121/128] Remove mention to flyteidl@union_type in doc-requirements Signed-off-by: Eduardo Apolinario --- .github/workflows/pythonbuild.yml | 1 - 1 file changed, 1 deletion(-) diff --git a/.github/workflows/pythonbuild.yml b/.github/workflows/pythonbuild.yml index 4d6db42256..289fa802a5 100644 --- a/.github/workflows/pythonbuild.yml +++ b/.github/workflows/pythonbuild.yml @@ -154,6 +154,5 @@ jobs: run: | python -m pip install --upgrade pip==21.2.4 setuptools wheel pip install -r doc-requirements.txt - git clone https://github.com/flyteorg/flyteidl.git && cd flyteidl && git checkout union_type && pip install . && cd .. - name: Build the documentation run: make -C docs html From 843cda99f121aba2954623f7f1946e31193edc39 Mon Sep 17 00:00:00 2001 From: Eduardo Apolinario Date: Thu, 17 Mar 2022 17:34:49 -0700 Subject: [PATCH 122/128] Linting Signed-off-by: Eduardo Apolinario --- tests/flytekit/unit/core/test_type_engine.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 4967f7cda5..5a53c11d19 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -1145,7 +1145,7 @@ def test_nested_annotated(): """ Test to show that nested Annotated types are flattened. """ - pt = Annotated[Annotated[int, 'inner-annotation'], 'outer-annotation'] + pt = Annotated[Annotated[int, "inner-annotation"], "outer-annotation"] lt = TypeEngine.to_literal_type(pt) assert lt.simple == model_types.SimpleType.INTEGER From 861b157aa27cf9dc3b2233cde86a0e78b9fd6c45 Mon Sep 17 00:00:00 2001 From: Eduardo Apolinario Date: Thu, 17 Mar 2022 18:25:04 -0700 Subject: [PATCH 123/128] Use tempfile.mkdtemp to create a temporary directory for local data persistence. Signed-off-by: Eduardo Apolinario --- flytekit/core/data_persistence.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/flytekit/core/data_persistence.py b/flytekit/core/data_persistence.py index cd4ffd1619..e24febefef 100644 --- a/flytekit/core/data_persistence.py +++ b/flytekit/core/data_persistence.py @@ -22,10 +22,10 @@ """ -import datetime import os import pathlib import re +import tempfile import typing from abc import abstractmethod from distutils import dir_util @@ -448,11 +448,9 @@ def put_data(self, local_path: Union[str, os.PathLike], remote_path: str, is_mul DataPersistencePlugins.register_plugin("file://", DiskPersistence) DataPersistencePlugins.register_plugin("/", DiskPersistence) -tmp_dir_prefix = f"{os.sep}tmp{os.sep}flyte" - -tmp_dir = os.path.join(tmp_dir_prefix, datetime.datetime.now().strftime("%Y%m%d_%H%M%S")) +tmp_dir_prefix = tempfile.mkdtemp(prefix="tmp-flyte") default_local_file_access_provider = FileAccessProvider( - local_sandbox_dir=os.path.join(tmp_dir, "sandbox"), - raw_output_prefix=os.path.join(tmp_dir, "raw"), + local_sandbox_dir=os.path.join(tmp_dir_prefix, "sandbox"), + raw_output_prefix=os.path.join(tmp_dir_prefix, "raw"), data_config=DataConfig.auto(), ) From 2657c2560059b7289cd0f61425302f08c57a5993 Mon Sep 17 00:00:00 2001 From: eduardo apolinario Date: Thu, 17 Mar 2022 21:06:56 -0700 Subject: [PATCH 124/128] Revert "Use tempfile.mkdtemp to create a temporary directory for local data persistence." This reverts commit 861b157aa27cf9dc3b2233cde86a0e78b9fd6c45. Signed-off-by: eduardo apolinario --- flytekit/core/data_persistence.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/flytekit/core/data_persistence.py b/flytekit/core/data_persistence.py index e24febefef..cd4ffd1619 100644 --- a/flytekit/core/data_persistence.py +++ b/flytekit/core/data_persistence.py @@ -22,10 +22,10 @@ """ +import datetime import os import pathlib import re -import tempfile import typing from abc import abstractmethod from distutils import dir_util @@ -448,9 +448,11 @@ def put_data(self, local_path: Union[str, os.PathLike], remote_path: str, is_mul DataPersistencePlugins.register_plugin("file://", DiskPersistence) DataPersistencePlugins.register_plugin("/", DiskPersistence) -tmp_dir_prefix = tempfile.mkdtemp(prefix="tmp-flyte") +tmp_dir_prefix = f"{os.sep}tmp{os.sep}flyte" + +tmp_dir = os.path.join(tmp_dir_prefix, datetime.datetime.now().strftime("%Y%m%d_%H%M%S")) default_local_file_access_provider = FileAccessProvider( - local_sandbox_dir=os.path.join(tmp_dir_prefix, "sandbox"), - raw_output_prefix=os.path.join(tmp_dir_prefix, "raw"), + local_sandbox_dir=os.path.join(tmp_dir, "sandbox"), + raw_output_prefix=os.path.join(tmp_dir, "raw"), data_config=DataConfig.auto(), ) From 1672bf43867e1fff05531eaf03eada52348583ef Mon Sep 17 00:00:00 2001 From: eduardo apolinario Date: Thu, 17 Mar 2022 21:12:16 -0700 Subject: [PATCH 125/128] Force temporary file to not be deleted in test Signed-off-by: eduardo apolinario --- tests/flytekit/unit/core/test_type_hints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/flytekit/unit/core/test_type_hints.py b/tests/flytekit/unit/core/test_type_hints.py index 240a66976a..ee4bc72756 100644 --- a/tests/flytekit/unit/core/test_type_hints.py +++ b/tests/flytekit/unit/core/test_type_hints.py @@ -1621,7 +1621,7 @@ def wf(a: ut) -> ut: assert wf(a=2) == 2 assert wf(a="2") == "2" assert wf(a=2.0) == 2.0 - file = tempfile.NamedTemporaryFile() + file = tempfile.NamedTemporaryFile(delete=False) assert isinstance(wf(a=FlyteFile(file.name)), FlyteFile) assert isinstance(wf(a=FlyteSchema()), FlyteSchema) assert wf(a=[1, 2, 3]) == [1, 2, 3] From 7349bcf4d4cab6ae5ff749abfddfe1203b4b74f0 Mon Sep 17 00:00:00 2001 From: eduardo apolinario Date: Thu, 17 Mar 2022 21:33:38 -0700 Subject: [PATCH 126/128] Regenerate papermill dev-requirements Signed-off-by: eduardo apolinario --- plugins/flytekit-papermill/dev-requirements.in | 2 +- plugins/flytekit-papermill/dev-requirements.txt | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/plugins/flytekit-papermill/dev-requirements.in b/plugins/flytekit-papermill/dev-requirements.in index f6a875fa41..c056eb983b 100644 --- a/plugins/flytekit-papermill/dev-requirements.in +++ b/plugins/flytekit-papermill/dev-requirements.in @@ -1,3 +1,3 @@ -flyteidl>=0.23.0 +flyteidl>=0.24.0 git+https://github.com/flyteorg/flytekit@master#egg=flytekitplugins-spark&subdirectory=plugins/flytekit-spark # vcs+protocol://repo_url/#egg=pkg&subdirectory=flyte diff --git a/plugins/flytekit-papermill/dev-requirements.txt b/plugins/flytekit-papermill/dev-requirements.txt index 54c6aff4cd..5e7a446b3b 100644 --- a/plugins/flytekit-papermill/dev-requirements.txt +++ b/plugins/flytekit-papermill/dev-requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.8 +# This file is autogenerated by pip-compile with python 3.9 # To update, run: # # pip-compile dev-requirements.in @@ -38,13 +38,13 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.13 # via flytekit -flyteidl==0.23.0 +flyteidl==0.24.0 # via # -r dev-requirements.in # flytekit flytekit==0.30.0 # via flytekitplugins-spark -git+https://github.com/flyteorg/flytekit@master#egg=flytekitplugins-spark&subdirectory=plugins/flytekit-spark +flytekitplugins-spark @ git+https://github.com/flyteorg/flytekit@master#subdirectory=plugins/flytekit-spark # via -r dev-requirements.in googleapis-common-protos==1.55.0 # via flyteidl From 8271d2cb51fa7936ad565b84701d4cab3b51563b Mon Sep 17 00:00:00 2001 From: Eduardo Apolinario Date: Mon, 21 Mar 2022 11:03:30 -0700 Subject: [PATCH 127/128] Remove duplicate code Signed-off-by: Eduardo Apolinario --- tests/flytekit/unit/core/test_type_engine.py | 16 ---------------- 1 file changed, 16 deletions(-) diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 5a53c11d19..cbfa7f6d03 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -382,22 +382,6 @@ def test_guessing_basic(): pt = TypeEngine.guess_python_type(lt) assert pt is FlytePickle - lt = model_types.LiteralType( - blob=BlobType( - format=FlytePickleTransformer.PYTHON_PICKLE_FORMAT, dimensionality=BlobType.BlobDimensionality.SINGLE - ) - ) - pt = TypeEngine.guess_python_type(lt) - assert pt is FlytePickle - - lt = model_types.LiteralType( - blob=BlobType( - format=FlytePickleTransformer.PYTHON_PICKLE_FORMAT, dimensionality=BlobType.BlobDimensionality.SINGLE - ) - ) - pt = TypeEngine.guess_python_type(lt) - assert pt is FlytePickle - def test_guessing_containers(): b = model_types.LiteralType(simple=model_types.SimpleType.BOOLEAN) From b94f1475b8baef594df8c57fbb345eeb76451692 Mon Sep 17 00:00:00 2001 From: Eduardo Apolinario Date: Mon, 21 Mar 2022 11:35:47 -0700 Subject: [PATCH 128/128] Put a lower bound on the pip version installed in CI Signed-off-by: Eduardo Apolinario --- Makefile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Makefile b/Makefile index bd14c63aa5..d5d41a4af2 100644 --- a/Makefile +++ b/Makefile @@ -12,7 +12,7 @@ help: .PHONY: install-piptools install-piptools: - pip install -U pip-tools setuptools wheel pip==22.0.3 + pip install -U pip-tools setuptools wheel "pip>=22.0.3" .PHONY: update_boilerplate update_boilerplate: