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

flyte/remote/_client/_protocols.py

@@ -58,6 +58,8 @@ class TaskService(Protocol):
         self, request: task_service_pb2.GetTaskDetailsRequest
     ) -> task_service_pb2.GetTaskDetailsResponse: ...
 
+    async def ListTasks(self, request: task_service_pb2.ListTasksRequest) -> task_service_pb2.ListTasksResponse: ...
+
 
 class RunService(Protocol):
     async def CreateRun(self, request: run_service_pb2.CreateRunRequest) -> run_service_pb2.CreateRunResponse: ...

flyte/remote/_task.py

@@ -3,27 +3,44 @@ from __future__ import annotations
 import functools
 from dataclasses import dataclass
 from threading import Lock
-from typing import Any, Callable, Coroutine, Dict, Literal, Optional, Union
+from typing import Any, AsyncIterator, Callable, Coroutine, Dict, Iterator, Literal, Optional, Tuple, Union
 
 import rich.repr
+from google.protobuf import timestamp
 
 import flyte
 import flyte.errors
 from flyte._context import internal_ctx
-from flyte._initialize import get_client, get_common_config
+from flyte._initialize import ensure_client, get_client, get_common_config
+from flyte._logging import logger
+from flyte._protos.common import list_pb2
 from flyte._protos.workflow import task_definition_pb2, task_service_pb2
 from flyte.models import NativeInterface
 from flyte.syncify import syncify
 
 
+def _repr_task_metadata(metadata: task_definition_pb2.TaskMetadata) -> rich.repr.Result:
+    """
+    Rich representation of the task metadata.
+    """
+    if metadata.deployed_by:
+        if metadata.deployed_by.user:
+            yield "deployed_by", f"User: {metadata.deployed_by.user.spec.email}"
+        else:
+            yield "deployed_by", f"App: {metadata.deployed_by.application.spec.name}"
+    yield "short_name", metadata.short_name
+    yield "deployed_at", timestamp.to_datetime(metadata.deployed_at)
+    yield "environment_name", metadata.environment_name
+
+
 class LazyEntity:
     """
     Fetches the entity when the entity is called or when the entity is retrieved.
     The entity is derived from RemoteEntity so that it behaves exactly like the mimicked entity.
     """
 
-    def __init__(self, name: str, getter: Callable[..., Coroutine[Any, Any, Task]], *args, **kwargs):
-        self._task: Optional[Task] = None
+    def __init__(self, name: str, getter: Callable[..., Coroutine[Any, Any, TaskDetails]], *args, **kwargs):
+        self._task: Optional[TaskDetails] = None
         self._getter = getter
         self._name = name
         self._mutex = Lock()
@@ -33,7 +50,7 @@ class LazyEntity:
         return self._name
 
     @syncify
-    async def fetch(self) -> Task:
+    async def fetch(self) -> TaskDetails:
         """
         Forwards all other attributes to task, causing the task to be fetched!
         """
@@ -62,7 +79,7 @@ AutoVersioning = Literal["latest", "current"]
 
 
 @dataclass
-class Task:
+class TaskDetails:
     pb2: task_definition_pb2.TaskDetails
 
     @classmethod
@@ -87,10 +104,19 @@ class Task:
         if version is None and auto_version not in ["latest", "current"]:
             raise ValueError("auto_version must be either 'latest' or 'current'.")
 
-        async def deferred_get(_version: str | None, _auto_version: AutoVersioning | None) -> Task:
+        async def deferred_get(_version: str | None, _auto_version: AutoVersioning | None) -> TaskDetails:
             if _version is None:
                 if _auto_version == "latest":
-                    raise NotImplementedError("auto_version=latest is not yet implemented.")
+                    tasks = []
+                    async for x in Task.listall.aio(
+                        by_task_name=name,
+                        sort_by=("created_at", "desc"),
+                        limit=1,
+                    ):
+                        tasks.append(x)
+                    if not tasks:
+                        raise flyte.errors.ReferenceTaskError(f"Task {name} not found.")
+                    _version = tasks[0].version
                 elif _auto_version == "current":
                     ctx = flyte.ctx()
                     if ctx is None:
@@ -205,7 +231,7 @@ class Task:
         env: Optional[Dict[str, str]] = None,
         secrets: Optional[flyte.SecretRequest] = None,
         **kwargs: Any,
-    ) -> Task:
+    ) -> TaskDetails:
         raise NotImplementedError
 
     def __rich_repr__(self) -> rich.repr.Result:
@@ -223,5 +249,111 @@ class Task:
         yield "resources", self.resources
 
 
+@dataclass
+class Task:
+    pb2: task_definition_pb2.Task
+
+    def __init__(self, pb2: task_definition_pb2.Task):
+        self.pb2 = pb2
+
+    @property
+    def name(self) -> str:
+        """
+        The name of the task.
+        """
+        return self.pb2.task_id.name
+
+    @property
+    def version(self) -> str:
+        """
+        The version of the task.
+        """
+        return self.pb2.task_id.version
+
+    @classmethod
+    def get(cls, name: str, version: str | None = None, auto_version: AutoVersioning | None = None) -> LazyEntity:
+        """
+        Get a task by its ID or name. If both are provided, the ID will take precedence.
+
+        Either version or auto_version are required parameters.
+
+        :param name: The name of the task.
+        :param version: The version of the task.
+        :param auto_version: If set to "latest", the latest-by-time ordered from now, version of the task will be used.
+         If set to "current", the version will be derived from the callee tasks context. This is useful if you are
+         deploying all environments with the same version. If auto_version is current, you can only access the task from
+         within a task context.
+        """
+        return TaskDetails.get(name, version=version, auto_version=auto_version)
+
+    @syncify
+    @classmethod
+    async def listall(
+        cls,
+        by_task_name: str | None = None,
+        sort_by: Tuple[str, Literal["asc", "desc"]] | None = None,
+        limit: int = 100,
+    ) -> Union[AsyncIterator[Task], Iterator[Task]]:
+        """
+        Get all runs for the current project and domain.
+
+        :param by_task_name: If provided, only tasks with this name will be returned.
+        :param sort_by: The sorting criteria for the project list, in the format (field, order).
+        :param limit: The maximum number of tasks to return.
+        :return: An iterator of runs.
+        """
+        ensure_client()
+        token = None
+        sort_by = sort_by or ("created_at", "asc")
+        sort_pb2 = list_pb2.Sort(
+            key=sort_by[0], direction=list_pb2.Sort.ASCENDING if sort_by[1] == "asc" else list_pb2.Sort.DESCENDING
+        )
+        cfg = get_common_config()
+        filters = []
+        if by_task_name:
+            filters.append(
+                list_pb2.Filter(
+                    function=list_pb2.Filter.Function.EQUAL,
+                    field="name",
+                    values=[by_task_name],
+                )
+            )
+        original_limit = limit
+        if limit > cfg.batch_size:
+            limit = cfg.batch_size
+        retrieved = 0
+        while True:
+            resp = await get_client().task_service.ListTasks(
+                task_service_pb2.ListTasksRequest(
+                    org=cfg.org,
+                    request=list_pb2.ListRequest(
+                        sort_by=sort_pb2,
+                        filters=filters,
+                        limit=limit,
+                        token=token,
+                    ),
+                )
+            )
+            token = resp.token
+            for t in resp.tasks:
+                retrieved += 1
+                yield cls(t)
+            if not token or retrieved >= original_limit:
+                logger.debug(f"Retrieved {retrieved} tasks, stopping iteration.")
+                break
+
+    def __rich_repr__(self) -> rich.repr.Result:
+        """
+        Rich representation of the task.
+        """
+        yield "project", self.pb2.task_id.project
+        yield "domain", self.pb2.task_id.domain
+        yield "name", self.pb2.task_id.name
+        yield "version", self.pb2.task_id.version
+        yield "short_name", self.pb2.metadata.short_name
+        for t in _repr_task_metadata(self.pb2.metadata):
+            yield t
+
+
 if __name__ == "__main__":
     tk = Task.get(name="example_task")

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.
     """
 
@@ -141,7 +141,6 @@ class _BackgroundLoop:
                 aio_future: asyncio.Future[R_co] = asyncio.wrap_future(future)
                 # await for the future to complete and yield its result
                 v = await aio_future
-                print(f"Yielding value: {v}")
                 yield v
             except StopAsyncIteration:
                 break

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.