flyte 0.0.1b3__py3-none-any.whl → 0.2.0a0__py3-none-any.whl

flyte/types/_type_engine.py

@@ -35,13 +35,14 @@ from mashumaro.jsonschema.models import Context, JSONSchema
 from mashumaro.jsonschema.plugins import BasePlugin
 from mashumaro.jsonschema.schema import Instance
 from mashumaro.mixins.json import DataClassJSONMixin
+from pydantic import BaseModel
 from typing_extensions import Annotated, get_args, get_origin
 
 import flyte.storage as storage
-from flyte._datastructures import NativeInterface
 from flyte._hash import HashMethod
 from flyte._logging import logger
 from flyte._utils.helpers import load_proto_from_file
+from flyte.models import NativeInterface
 
 from ._utils import literal_types_match
 
@@ -352,6 +353,79 @@ class RestrictedTypeTransformer(TypeTransformer[T], ABC):
         raise RestrictedTypeError(f"Transformer for type {self.python_type} is restricted currently")
 
 
+class PydanticTransformer(TypeTransformer[BaseModel]):
+    def __init__(self):
+        super().__init__("Pydantic Transformer", BaseModel, enable_type_assertions=False)
+
+    def get_literal_type(self, t: Type[BaseModel]) -> LiteralType:
+        schema = t.model_json_schema()
+        fields = t.__annotations__.items()
+
+        literal_type = {}
+        for name, python_type in fields:
+            try:
+                literal_type[name] = TypeEngine.to_literal_type(python_type)
+            except Exception as e:
+                logger.warning(
+                    "Field {} of type {} cannot be converted to a literal type. Error: {}".format(name, python_type, e)
+                )
+
+        # This is for attribute access in FlytePropeller.
+        ts = TypeStructure(tag="", dataclass_type=literal_type)
+
+        meta_struct = struct_pb2.Struct()
+        meta_struct.update(
+            {
+                CACHE_KEY_METADATA: {
+                    SERIALIZATION_FORMAT: MESSAGEPACK,
+                }
+            }
+        )
+
+        return LiteralType(
+            simple=SimpleType.STRUCT,
+            metadata=schema,
+            structure=ts,
+            annotation=TypeAnnotation(annotations=meta_struct),
+        )
+
+    async def to_literal(
+        self,
+        python_val: BaseModel,
+        python_type: Type[BaseModel],
+        expected: LiteralType,
+    ) -> Literal:
+        json_str = python_val.model_dump_json()
+        dict_obj = json.loads(json_str)
+        msgpack_bytes = msgpack.dumps(dict_obj)
+        return Literal(scalar=Scalar(binary=Binary(value=msgpack_bytes, tag=MESSAGEPACK)))
+
+    def from_binary_idl(self, binary_idl_object: Binary, expected_python_type: Type[BaseModel]) -> BaseModel:
+        if binary_idl_object.tag == MESSAGEPACK:
+            dict_obj = msgpack.loads(binary_idl_object.value, strict_map_key=False)
+            json_str = json.dumps(dict_obj)
+            python_val = expected_python_type.model_validate_json(
+                json_data=json_str, strict=False, context={"deserialize": True}
+            )
+            return python_val
+        else:
+            raise TypeTransformerFailedError(f"Unsupported binary format: `{binary_idl_object.tag}`")
+
+    async def to_python_value(self, lv: Literal, expected_python_type: Type[BaseModel]) -> BaseModel:
+        """
+        There are two kinds of literal values to handle:
+        1. Protobuf Structs (from the UI)
+        2. Binary scalars (from other sources)
+        We need to account for both cases accordingly.
+        """
+        if lv and lv.HasField("scalar") and lv.scalar.HasField("binary"):
+            return self.from_binary_idl(lv.scalar.binary, expected_python_type)  # type: ignore
+
+        json_str = _json_format.MessageToJson(lv.scalar.generic)
+        python_val = expected_python_type.model_validate_json(json_str, strict=False, context={"deserialize": True})
+        return python_val
+
+
 class PydanticSchemaPlugin(BasePlugin):
     """This allows us to generate proper schemas for Pydantic models."""
 
@@ -450,7 +524,6 @@ class DataclassTransformer(TypeTransformer[object]):
         # However, FooSchema is created by flytekit and it's not equal to the user-defined dataclass (Foo).
         # Therefore, we should iterate all attributes in the dataclass and check the type of value in dataclass
         #   matches the expected_type.
-
         expected_fields_dict = {}
 
         for f in dataclasses.fields(expected_type):
@@ -563,9 +636,8 @@ class DataclassTransformer(TypeTransformer[object]):
 
         # This is for attribute access in FlytePropeller.
         ts = TypeStructure(tag="", dataclass_type=literal_type)
-        from google.protobuf.struct_pb2 import Struct
 
-        meta_struct = Struct()
+        meta_struct = struct_pb2.Struct()
         meta_struct.update(
             {
                 CACHE_KEY_METADATA: {
@@ -628,7 +700,7 @@ class DataclassTransformer(TypeTransformer[object]):
                 field.type = self._get_origin_type_in_annotation(cast(type, field.type))
         return python_type
 
-    async def from_binary_idl(self, binary_idl_object: Binary, expected_python_type: Type[T]) -> T:
+    def from_binary_idl(self, binary_idl_object: Binary, expected_python_type: Type[T]) -> T:
         if binary_idl_object.tag == MESSAGEPACK:
             if issubclass(expected_python_type, DataClassJSONMixin):
                 dict_obj = msgpack.loads(binary_idl_object.value, strict_map_key=False)
@@ -653,9 +725,10 @@ class DataclassTransformer(TypeTransformer[object]):
                 "user defined datatypes in Flytekit"
             )
 
-        if lv.scalar and lv.scalar.binary:
-            return await self.from_binary_idl(lv.scalar.binary, expected_python_type)  # type: ignore
+        if lv.HasField("scalar") and lv.scalar.HasField("binary"):
+            return self.from_binary_idl(lv.scalar.binary, expected_python_type)  # type: ignore
 
+        # todo: revisit this, it should always be a binary in v2.
         json_str = _json_format.MessageToJson(lv.scalar.generic)
 
         # The `from_json` function is provided from mashumaro's `DataClassJSONMixin`.
@@ -971,11 +1044,10 @@ class TypeEngine(typing.Generic[T]):
                 return cls._REGISTRY[python_type.__origin__]
 
         # Handling UnionType specially - PEP 604
-        if sys.version_info >= (3, 10):
-            import types
+        import types
 
-            if isinstance(python_type, types.UnionType):
-                return cls._REGISTRY[types.UnionType]
+        if isinstance(python_type, types.UnionType):
+            return cls._REGISTRY[types.UnionType]
 
         if python_type in cls._REGISTRY:
             return cls._REGISTRY[python_type]
@@ -1009,7 +1081,7 @@ class TypeEngine(typing.Generic[T]):
             return cls._DATACLASS_TRANSFORMER
 
         display_pickle_warning(str(python_type))
-        from flyte.io.pickle.transformer import FlytePickleTransformer
+        from flyte.types._pickle import FlytePickleTransformer
 
         return FlytePickleTransformer()
 
@@ -1022,7 +1094,7 @@ class TypeEngine(typing.Generic[T]):
             # Avoid a race condition where concurrent threads may exit lazy_import_transformers before the transformers
             # have been imported. This could be implemented without a lock if you assume python assignments are atomic
             # and re-registering transformers is acceptable, but I decided to play it safe.
-            from flyte.io.structured_dataset import lazy_import_structured_dataset_handler
+            from flyte.io import lazy_import_structured_dataset_handler
 
             # todo: bring in extras transformers (pytorch, etc.)
             lazy_import_structured_dataset_handler()
@@ -1208,7 +1280,9 @@ class TypeEngine(typing.Generic[T]):
                 python_type = type_hints.get(k, type(d[k]))
                 e: BaseException = literal_map[k].exception()  # type: ignore
                 if isinstance(e, TypeError):
-                    raise TypeError(f"Error converting: {type(v)}, {python_type}, received_value {v}")
+                    raise TypeError(
+                        f"Error converting: Var:{k}, type:{type(v)}, into:{python_type}, received_value {v}"
+                    )
                 else:
                     raise e
             literal_map[k] = v.result()
@@ -1658,7 +1732,7 @@ class DictTransformer(TypeTransformer[dict]):
         Converts a Python dictionary to a Flyte-specific ``Literal`` using MessagePack encoding.
         Falls back to Pickle if encoding fails and `allow_pickle` is True.
         """
-        from flyte.io.pickle.transformer import FlytePickle
+        from flyte.types._pickle import FlytePickle
 
         try:
             # Handle dictionaries with non-string keys (e.g., Dict[int, Type])
@@ -1764,7 +1838,7 @@ class DictTransformer(TypeTransformer[dict]):
         # pr: han-ru is this part still necessary?
         if lv and lv.HasField("scalar") and lv.scalar.HasField("generic"):
             if lv.metadata and lv.metadata.get("format", None) == "pickle":
-                from flyte.io.pickle.transformer import FlytePickle
+                from flyte.types._pickle import FlytePickle
 
                 uri = json.loads(_json_format.MessageToJson(lv.scalar.generic)).get("pickle_file")
                 return await FlytePickle.from_pickle(uri)
@@ -1992,7 +2066,9 @@ DatetimeTransformer = SimpleTransformer(
     datetime.datetime,
     types_pb2.LiteralType(simple=types_pb2.SimpleType.DATETIME),
     lambda x: Literal(scalar=Scalar(primitive=Primitive(datetime=x))),
-    lambda x: x.scalar.primitive.datetime if x.scalar.primitive.HasField("datetime") else None,
+    lambda x: x.scalar.primitive.datetime.ToDatetime().replace(tzinfo=datetime.timezone.utc)
+    if x.scalar.primitive.HasField("datetime")
+    else None,
 )
 
 TimedeltaTransformer = SimpleTransformer(
@@ -2000,7 +2076,7 @@ TimedeltaTransformer = SimpleTransformer(
     datetime.timedelta,
     types_pb2.LiteralType(simple=types_pb2.SimpleType.DURATION),
     lambda x: Literal(scalar=Scalar(primitive=Primitive(duration=x))),
-    lambda x: x.scalar.primitive.duration if x.scalar.primitive.HasField("duration") else None,
+    lambda x: x.scalar.primitive.duration.ToTimedelta() if x.scalar.primitive.HasField("duration") else None,
 )
 
 DateTransformer = SimpleTransformer(
@@ -2038,6 +2114,7 @@ def _register_default_type_transformers():
     TypeEngine.register(DictTransformer())
     TypeEngine.register(EnumTransformer())
     TypeEngine.register(ProtobufTransformer())
+    TypeEngine.register(PydanticTransformer())
 
     # inner type is. Also unsupported are typing's Tuples. Even though you can look inside them, Flyte's type system
     # doesn't support these currently.

flyte-0.2.0a0.dist-info/METADATA

@@ -0,0 +1,249 @@
+Metadata-Version: 2.4
+Name: flyte
+Version: 0.2.0a0
+Summary: Add your description here
+Author-email: Ketan Umare <kumare3@users.noreply.github.com>
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: aiofiles>=24.1.0
+Requires-Dist: click>=8.2.1
+Requires-Dist: flyteidl==1.15.4b0
+Requires-Dist: cloudpickle>=3.1.1
+Requires-Dist: fsspec>=2025.3.0
+Requires-Dist: grpcio>=1.71.0
+Requires-Dist: obstore>=0.6.0
+Requires-Dist: protobuf>=6.30.1
+Requires-Dist: pydantic>=2.10.6
+Requires-Dist: pyyaml>=6.0.2
+Requires-Dist: rich-click>=1.8.9
+Requires-Dist: httpx>=0.28.1
+Requires-Dist: keyring>=25.6.0
+Requires-Dist: msgpack>=1.1.0
+Requires-Dist: toml>=0.10.2
+Requires-Dist: async-lru>=2.0.5
+Requires-Dist: mashumaro
+Requires-Dist: dataclasses_json
+
+# Flyte v2 SDK
+
+The next-generation SDK for Flyte.
+
+[![Publish Python Packages and Official Images](https://github.com/unionai/unionv2/actions/workflows/publish.yml/badge.svg)](https://github.com/unionai/unionv2/actions/workflows/publish.yml)
+
+## Quick start
+
+1. Run `uv venv`, and `source .venv/bin/activate` to create a new virtual environment.
+2. Install the latest version of the SDK by running the following:
+
+```
+uv pip install --no-cache --prerelease=allow --upgrade flyte
+```
+
+4. Create the config and point it to your cluster by running the following:
+
+```
+flyte create config --endpoint <your-endpoint-url> --project <your-project> --domain <your-domain>
+```
+
+This will create a `config.yaml` file in the current directory which will be referenced ahead of any other `config.yaml`s found in your system.
+
+5. Now you can run code with the CLI:
+
+```
+flyte run <path-to-your-script> <task-name>
+```
+
+## Hello World Example
+
+```python
+# hello_world.py
+
+import flyte
+
+env = flyte.TaskEnvironment(name="hello_world")
+
+
+@env.task
+async def say_hello(data: str) -> str:
+    return f"Hello {data}"
+
+
+@env.task
+async def say_hello_nested(data: str) -> str:
+    return await say_hello.override(resources=flyte.Resources(gpu="A100 80G:4")).execute(data)
+
+
+if __name__ == "__main__":
+    import asyncio
+
+    # to run pure python - the SDK is not invoked at all
+    asyncio.run(say_hello_nested("test"))
+
+    # To run locally, but run through type system etc
+    flyte.init()
+    flyte.run(say_hello_nested, "World")
+
+    # To run remote
+    flyte.init(endpoint="dns:///localhost:8090", insecure=True)
+    flyte.run(say_hello_nested, "World")
+    # It is possible to switch local and remote, but keeping init to have and endpoint, but , changing context during run
+    flyte.with_runcontext(mode="local").run(...)  # this will run locally only
+
+    # To run remote with a config
+    flyte.init_from_config("config.yaml")
+```
+
+## CLI
+
+All commands can be run from any root directory.
+For examples, it is not needed to have `__init__.py` in the directory.
+If you run from a directory, the code will automatically package and upload all modules that are imported.
+You can change the behavior by using `--copy-style` flag.
+
+```bash
+flyte run hello_world.py say_hello --data "World"
+```
+
+To follow the logs for the `a0` action, you can use the `--follow` flag:
+
+```bash
+flyte run --follow hello_world.py say_hello --data "World"
+```
+
+Note that `--follow` has to be used with the `run` command.
+
+Change copy style:
+
+```bash
+flyte run --copy-style all hello_world.py say_hello_nested --data "World"
+```
+
+## Building Images
+
+```python
+import flyte
+
+env = flyte.TaskEnvironment(
+    name="hello_world",
+    image=flyte.Image.auto().with_apt_packages(...).with_pip_packages(...),
+)
+
+```
+
+## Deploying
+
+```bash
+flyte deploy hello_world.py say_hello_nested
+```
+
+## Get information
+
+Get all runs:
+
+```bash
+flyte get run
+```
+
+Get a specific run:
+
+```bash
+flyte get run "run-name"
+```
+
+Get all actions for a run:
+
+```bash
+flyte get actions "run-name"
+```
+
+Get a specific action for a run:
+
+```bash
+flyte get action "run-name" "action-name"
+```
+
+Get action logs:
+
+```bash
+flyte get logs "run-name" ["action-name"]
+```
+
+This defaults to root action if no action name is provided
+
+## Running workflows programmatically in Python
+
+You can run any workflow programmatically within the script module using __main__:
+
+```python
+if __name__ == "__main__":
+    import flyte
+    flyte.init()
+    flyte.run(say_hello_nested, "World")
+```
+
+## Running scripts with dependencies specified in metadata headers
+
+You can also run a `uv` script with dependencies specified in metadata headers
+and build the task image automatically based on those dependencies:
+
+```python
+# container_images.py
+
+# /// script
+# dependencies = [
+#    "polars",
+#    "flyte>=0.2.0b12"
+# ]
+# ///
+
+import polars as pl
+
+import flyte
+
+
+env = flyte.TaskEnvironment(
+    name="polars_image",
+    image=flyte.Image.from_uv_script(
+        __file__,
+        name="flyte",
+        registry="ghcr.io/<you-username>"
+        arch=("linux/amd64", "linux/arm64"),
+    ).with_apt_packages("ca-certificates"),
+)
+
+
+@env.task
+async def create_dataframe() -> pl.DataFrame:
+    return pl.DataFrame(
+        {"name": ["Alice", "Bob", "Charlie"], "age": [25, 32, 37], "city": ["New York", "Paris", "Berlin"]}
+    )
+
+
+@env.task
+async def print_dataframe(dataframe: pl.DataFrame):
+    print(dataframe)
+
+
+@env.task
+async def workflow():
+    df = await create_dataframe()
+    await print_dataframe(df)
+
+
+if __name__ == "__main__":
+    flyte.init_from_config("config.yaml")
+    run = flyte.run(workflow)
+    print(run.name)
+    print(run.url)
+    run.wait(run)
+```
+
+When you execute
+
+```bash
+ uv run hello_world.py
+```
+
+`uv` will automatically update the local virtual environment with the dependencies specified in the metadata headers.
+Then, Flyte will build the task image using those dependencies and push it to the registry you specify.
+Flyte will then deploy the tasks to the cluster where the system will pull the image and run the tasks using it.