flyte 0.2.0b12__py3-none-any.whl → 0.2.0b13__py3-none-any.whl

flyte/_protos/workflow/state_service_pb2_grpc.py

@@ -15,15 +15,15 @@ class StateServiceStub(object):
         Args:
             channel: A grpc.Channel.
         """
-        self.Store = channel.unary_unary(
-                '/cloudidl.workflow.StateService/Store',
-                request_serializer=workflow_dot_state__service__pb2.StoreRequest.SerializeToString,
-                response_deserializer=workflow_dot_state__service__pb2.StoreResponse.FromString,
+        self.Put = channel.stream_stream(
+                '/cloudidl.workflow.StateService/Put',
+                request_serializer=workflow_dot_state__service__pb2.PutRequest.SerializeToString,
+                response_deserializer=workflow_dot_state__service__pb2.PutResponse.FromString,
                 )
-        self.Load = channel.unary_unary(
-                '/cloudidl.workflow.StateService/Load',
-                request_serializer=workflow_dot_state__service__pb2.LoadRequest.SerializeToString,
-                response_deserializer=workflow_dot_state__service__pb2.LoadResponse.FromString,
+        self.Get = channel.stream_stream(
+                '/cloudidl.workflow.StateService/Get',
+                request_serializer=workflow_dot_state__service__pb2.GetRequest.SerializeToString,
+                response_deserializer=workflow_dot_state__service__pb2.GetResponse.FromString,
                 )
         self.Watch = channel.unary_stream(
                 '/cloudidl.workflow.StateService/Watch',
@@ -36,15 +36,15 @@ class StateServiceServicer(object):
     """provides an interface for managing the state of actions.
     """
 
-    def Store(self, request, context):
-        """store the state of an action.
+    def Put(self, request_iterator, context):
+        """put the state of an action.
         """
         context.set_code(grpc.StatusCode.UNIMPLEMENTED)
         context.set_details('Method not implemented!')
         raise NotImplementedError('Method not implemented!')
 
-    def Load(self, request, context):
-        """load the state of an action.
+    def Get(self, request_iterator, context):
+        """get the state of an action.
         """
         context.set_code(grpc.StatusCode.UNIMPLEMENTED)
         context.set_details('Method not implemented!')
@@ -60,15 +60,15 @@ class StateServiceServicer(object):
 
 def add_StateServiceServicer_to_server(servicer, server):
     rpc_method_handlers = {
-            'Store': grpc.unary_unary_rpc_method_handler(
-                    servicer.Store,
-                    request_deserializer=workflow_dot_state__service__pb2.StoreRequest.FromString,
-                    response_serializer=workflow_dot_state__service__pb2.StoreResponse.SerializeToString,
+            'Put': grpc.stream_stream_rpc_method_handler(
+                    servicer.Put,
+                    request_deserializer=workflow_dot_state__service__pb2.PutRequest.FromString,
+                    response_serializer=workflow_dot_state__service__pb2.PutResponse.SerializeToString,
             ),
-            'Load': grpc.unary_unary_rpc_method_handler(
-                    servicer.Load,
-                    request_deserializer=workflow_dot_state__service__pb2.LoadRequest.FromString,
-                    response_serializer=workflow_dot_state__service__pb2.LoadResponse.SerializeToString,
+            'Get': grpc.stream_stream_rpc_method_handler(
+                    servicer.Get,
+                    request_deserializer=workflow_dot_state__service__pb2.GetRequest.FromString,
+                    response_serializer=workflow_dot_state__service__pb2.GetResponse.SerializeToString,
             ),
             'Watch': grpc.unary_stream_rpc_method_handler(
                     servicer.Watch,
@@ -87,7 +87,7 @@ class StateService(object):
     """
 
     @staticmethod
-    def Store(request,
+    def Put(request_iterator,
             target,
             options=(),
             channel_credentials=None,
@@ -97,14 +97,14 @@ class StateService(object):
             wait_for_ready=None,
             timeout=None,
             metadata=None):
-        return grpc.experimental.unary_unary(request, target, '/cloudidl.workflow.StateService/Store',
-            workflow_dot_state__service__pb2.StoreRequest.SerializeToString,
-            workflow_dot_state__service__pb2.StoreResponse.FromString,
+        return grpc.experimental.stream_stream(request_iterator, target, '/cloudidl.workflow.StateService/Put',
+            workflow_dot_state__service__pb2.PutRequest.SerializeToString,
+            workflow_dot_state__service__pb2.PutResponse.FromString,
             options, channel_credentials,
             insecure, call_credentials, compression, wait_for_ready, timeout, metadata)
 
     @staticmethod
-    def Load(request,
+    def Get(request_iterator,
             target,
             options=(),
             channel_credentials=None,
@@ -114,9 +114,9 @@ class StateService(object):
             wait_for_ready=None,
             timeout=None,
             metadata=None):
-        return grpc.experimental.unary_unary(request, target, '/cloudidl.workflow.StateService/Load',
-            workflow_dot_state__service__pb2.LoadRequest.SerializeToString,
-            workflow_dot_state__service__pb2.LoadResponse.FromString,
+        return grpc.experimental.stream_stream(request_iterator, target, '/cloudidl.workflow.StateService/Get',
+            workflow_dot_state__service__pb2.GetRequest.SerializeToString,
+            workflow_dot_state__service__pb2.GetResponse.FromString,
             options, channel_credentials,
             insecure, call_credentials, compression, wait_for_ready, timeout, metadata)
 

flyte/_run.py

@@ -63,6 +63,7 @@ class _Runner:
         raw_data_path: str | None = None,
         metadata_path: str | None = None,
         run_base_dir: str | None = None,
+        overwrite_cache: bool = False,
     ):
         init_config = _get_init_config()
         client = init_config.client if init_config else None
@@ -81,6 +82,7 @@ class _Runner:
         self._raw_data_path = raw_data_path
         self._metadata_path = metadata_path or "/tmp"
         self._run_base_dir = run_base_dir or "/tmp/base"
+        self._overwrite_cache = overwrite_cache
 
     @requires_initialization
     async def _run_remote(self, obj: TaskTemplate[P, R] | LazyEntity, *args: P.args, **kwargs: P.kwargs) -> Run:
@@ -182,6 +184,9 @@ class _Runner:
                         project_id=project_id,
                         task_spec=task_spec,
                         inputs=inputs.proto_inputs,
+                        run_spec=run_definition_pb2.RunSpec(
+                            overwrite_cache=self._overwrite_cache,
+                        ),
                     ),
                 )
                 return Run(pb2=resp.run)
@@ -414,6 +419,7 @@ def with_runcontext(
     interactive_mode: bool | None = None,
     raw_data_path: str | None = None,
     run_base_dir: str | None = None,
+    overwrite_cache: bool = False,
 ) -> _Runner:
     """
     Launch a new run with the given parameters as the context.

flyte/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.0b12'
-__version_tuple__ = version_tuple = (0, 2, 0, 'b12')
+__version__ = version = '0.2.0b13'
+__version_tuple__ = version_tuple = (0, 2, 0, 'b13')

flyte/cli/_deploy.py

@@ -147,6 +147,6 @@ deploy = EnvFiles(
     name="deploy",
     help="""
     Deploy one or more environments from a python file.
-    The deploy command will create or update environments in the Flyte system.
+    This command will create or update environments in the Flyte system.
     """,
 )

flyte/cli/_get.py

@@ -20,7 +20,7 @@ def get():
     Using a `get` subcommand without any arguments will retrieve a list of available resources to get.
     For example:
 
-    * `get project` (without specifiying aproject), will list all projects.
+    * `get project` (without specifying a project), will list all projects.
     * `get project my_project` will return the details of the project named `my_project`.
 
     In some cases, a partially specified command will act as a filter and return available further parameters.
@@ -143,7 +143,7 @@ def action(
     "--pretty",
     is_flag=True,
     default=False,
-    help="Show logs in a auto scrolling box, where number of lines is limited to `--lines`",
+    help="Show logs in an auto-scrolling box, where number of lines is limited to `--lines`",
 )
 @click.option(
     "--attempt", "-a", type=int, default=None, help="Attempt number to show logs for, defaults to the latest attempt."

flyte/cli/_run.py

@@ -66,7 +66,7 @@ class RunArguments:
                 ["--follow", "-f"],
                 is_flag=True,
                 default=False,
-                help="Wait and watch logs for the parent action. If not provided, the cli will exit after "
+                help="Wait and watch logs for the parent action. If not provided, the CLI will exit after "
                 "successfully launching a remote execution with a link to the UI.",
             )
         },
@@ -108,7 +108,6 @@ class RunTaskCommand(click.Command):
 
             r = flyte.with_runcontext(
                 copy_style=self.run_args.copy_style,
-                version=self.run_args.copy_style,
                 mode="local" if self.run_args.local else "remote",
                 name=self.run_args.name,
             ).run(self.obj, **ctx.params)

flyte/cli/main.py

@@ -103,7 +103,7 @@ def main(
     config_file: str | None,
 ):
     """
-    The Flyte CLI is the the command line interface for working with the Flyte SDK and backend.
+    The Flyte CLI is the command line interface for working with the Flyte SDK and backend.
 
     It follows a simple verb/noun structure,
     where the top-level commands are verbs that describe the action to be taken,

flyte/syncify/_api.py

@@ -50,7 +50,7 @@ class SyncGenFunction(Protocol[P, R_co]):
 
 class _BackgroundLoop:
     """
-    A background event loop that runs in a separate thread and used the the Syncify decorator to run asynchronous
+    A background event loop that runs in a separate thread and used the `Syncify` decorator to run asynchronous
     functions or methods synchronously.
     """
 

flyte/types/_type_engine.py

@@ -35,6 +35,7 @@ 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
@@ -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."""
 
@@ -562,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: {
@@ -627,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)
@@ -652,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`.
@@ -970,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]
@@ -2041,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.0b13.dist-info/METADATA

@@ -0,0 +1,249 @@
+Metadata-Version: 2.4
+Name: flyte
+Version: 0.2.0b13
+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.