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

{union/_cli → flyte/cli}/_params.py

@@ -1,4 +1,3 @@
-import asyncio
 import dataclasses
 import datetime
 import enum
@@ -7,6 +6,7 @@ import importlib.util
 import json
 import os
 import pathlib
+import re
 import sys
 import typing
 import typing as t
@@ -14,24 +14,16 @@ from typing import get_args
 
 import rich_click as click
 import yaml
+from click import Parameter
+from flyteidl.core.interface_pb2 import Variable
 from flyteidl.core.literals_pb2 import Literal
 from flyteidl.core.types_pb2 import BlobType, LiteralType, SimpleType
+from google.protobuf.json_format import MessageToDict
+from mashumaro.codecs.json import JSONEncoder
 
-from union._logging import logger
-from union.io import Dir, File
-from union.io.pickle.transformer import FlytePickleTransformer
-from union.storage._remote_fs import RemoteFSPathResolver
-from union.types import TypeEngine
-
-
-# ---------------------------------------------------
-# TODO replace these
-class ArtifactQuery:
-    pass
-
-
-def is_remote(v: str) -> bool:
-    return False
+from flyte._logging import logger
+from flyte.io import Dir, File
+from flyte.types._pickle import FlytePickleTransformer
 
 
 class StructuredDataset:
@@ -43,26 +35,6 @@ class StructuredDataset:
 # ---------------------------------------------------
 
 
-def is_pydantic_basemodel(python_type: typing.Type) -> bool:
-    """
-    Checks if the python type is a pydantic BaseModel
-    """
-    try:
-        import pydantic  # noqa: F401
-    except ImportError:
-        return False
-    else:
-        try:
-            from pydantic import BaseModel as BaseModelV2
-            from pydantic.v1 import BaseModel as BaseModelV1
-
-            return issubclass(python_type, BaseModelV1) or issubclass(python_type, BaseModelV2)
-        except ImportError:
-            from pydantic import BaseModel
-
-        return issubclass(python_type, BaseModel)
-
-
 def key_value_callback(_: typing.Any, param: str, values: typing.List[str]) -> typing.Optional[typing.Dict[str, str]]:
     """
     Callback for click to parse key-value pairs.
@@ -100,17 +72,13 @@ class DirParamType(click.ParamType):
     def convert(
         self, value: typing.Any, param: typing.Optional[click.Parameter], ctx: typing.Optional[click.Context]
     ) -> typing.Any:
-        if isinstance(value, ArtifactQuery):
-            return value
+        from flyte.storage import is_remote
 
-        # set remote_directory to false if running pyflyte run locally. This makes sure that the original
-        # directory is used and not a random one.
-        remote_directory = None if getattr(ctx.obj, "is_remote", False) else False
         if not is_remote(value):
             p = pathlib.Path(value)
             if not p.exists() or not p.is_dir():
                 raise click.BadParameter(f"parameter should be a valid flytedirectory path, {value}")
-        return Dir(path=value, remote_directory=remote_directory)
+        return Dir(path=value)
 
 
 class StructuredDatasetParamType(click.ParamType):
@@ -123,8 +91,6 @@ class StructuredDatasetParamType(click.ParamType):
     def convert(
         self, value: typing.Any, param: typing.Optional[click.Parameter], ctx: typing.Optional[click.Context]
     ) -> typing.Any:
-        if isinstance(value, ArtifactQuery):
-            return value
         if isinstance(value, str):
             return StructuredDataset(uri=value)
         elif isinstance(value, StructuredDataset):
@@ -138,22 +104,19 @@ class FileParamType(click.ParamType):
     def convert(
         self, value: typing.Any, param: typing.Optional[click.Parameter], ctx: typing.Optional[click.Context]
     ) -> typing.Any:
-        if isinstance(value, ArtifactQuery):
-            return value
-        # set remote_directory to false if running pyflyte run locally. This makes sure that the original
-        # file is used and not a random one.
-        remote_path = None if getattr(ctx.obj, "is_remote", False) else False
+        from flyte.storage import is_remote
+
         if not is_remote(value):
             p = pathlib.Path(value)
             if not p.exists() or not p.is_file():
                 raise click.BadParameter(f"parameter should be a valid file path, {value}")
-        return File(path=value, remote_path=remote_path)
+        return File.from_existing_remote(value)
 
 
 class PickleParamType(click.ParamType):
     name = "pickle"
 
-    def get_metavar(self, param: click.Parameter) -> t.Optional[str]:
+    def get_metavar(self, param: "Parameter", *args) -> t.Optional[str]:
         return "Python Object <Module>:<Object>"
 
     def convert(
@@ -163,7 +126,7 @@ class PickleParamType(click.ParamType):
             return value
         parts = value.split(":")
         if len(parts) != 2:
-            if ctx and ctx.obj and ctx.obj.verbose > 0:
+            if ctx and ctx.obj and ctx.obj.log_level >= 10:  # DEBUG level
                 click.echo(f"Did not receive a string in the expected format <MODULE>:<VAR>, falling back to: {value}")
             return value
         try:
@@ -185,9 +148,6 @@ class JSONIteratorParamType(click.ParamType):
         return value
 
 
-import re
-
-
 def parse_iso8601_duration(iso_duration: str) -> datetime.timedelta:
     pattern = re.compile(
         r"^P"  # Starts with 'P'
@@ -211,10 +171,10 @@ def parse_human_durations(text: str) -> list[datetime.timedelta]:
     durations = []
 
     for part in raw_parts:
-        part = part.strip().lower()
+        new_part = part.strip().lower()
 
         # Match 1:24 or :45
-        m_colon = re.match(r"^(?:(\d+):)?(\d+)$", part)
+        m_colon = re.match(r"^(?:(\d+):)?(\d+)$", new_part)
         if m_colon:
             minutes = int(m_colon.group(1)) if m_colon.group(1) else 0
             seconds = int(m_colon.group(2))
@@ -222,7 +182,7 @@ def parse_human_durations(text: str) -> list[datetime.timedelta]:
             continue
 
         # Match "10 days", "1 minute", etc.
-        m_units = re.match(r"^(\d+)\s*(day|hour|minute|second)s?$", part)
+        m_units = re.match(r"^(\d+)\s*(day|hour|minute|second)s?$", new_part)
         if m_units:
             value = int(m_units.group(1))
             unit = m_units.group(2)
@@ -270,9 +230,6 @@ class DateTimeType(click.DateTime):
     def convert(
         self, value: typing.Any, param: typing.Optional[click.Parameter], ctx: typing.Optional[click.Context]
     ) -> typing.Any:
-        if isinstance(value, ArtifactQuery):
-            return value
-
         if isinstance(value, str) and " " in value:
             import re
 
@@ -303,8 +260,6 @@ class DurationParamType(click.ParamType):
     def convert(
         self, value: typing.Any, param: typing.Optional[click.Parameter], ctx: typing.Optional[click.Context]
     ) -> typing.Any:
-        if isinstance(value, ArtifactQuery):
-            return value
         if value is None:
             raise click.BadParameter("None value cannot be converted to a Duration type.")
         return parse_duration(value)
@@ -318,8 +273,6 @@ class EnumParamType(click.Choice):
     def convert(
         self, value: typing.Any, param: typing.Optional[click.Parameter], ctx: typing.Optional[click.Context]
     ) -> enum.Enum:
-        if isinstance(value, ArtifactQuery):
-            return value
         if isinstance(value, self._enum_type):
             return value
         return self._enum_type(super().convert(value, param, ctx))
@@ -333,10 +286,7 @@ class UnionParamType(click.ParamType):
     def __init__(self, types: typing.List[click.ParamType]):
         super().__init__()
         self._types = self._sort_precedence(types)
-
-    @property
-    def name(self) -> str:
-        return "|".join([t.name for t in self._types])
+        self.name = "|".join([t.name for t in self._types])
 
     @staticmethod
     def _sort_precedence(tp: typing.List[click.ParamType]) -> typing.List[click.ParamType]:
@@ -350,7 +300,7 @@ class UnionParamType(click.ParamType):
                 str_types.append(p)
             else:
                 others.append(p)
-        return others + str_types + unprocessed
+        return others + str_types + unprocessed  # type: ignore
 
     def convert(
         self, value: typing.Any, param: typing.Optional[click.Parameter], ctx: typing.Optional[click.Context]
@@ -359,8 +309,6 @@ class UnionParamType(click.ParamType):
         Important to implement NoneType / Optional.
         Also could we just determine the click types from the python types
         """
-        if isinstance(value, ArtifactQuery):
-            return value
         for p in self._types:
             try:
                 return p.convert(value, param, ctx)
@@ -386,7 +334,7 @@ class JsonParamType(click.ParamType):
                 # We failed to load the json, so we'll try to load it as a file
                 if os.path.exists(value):
                     # if the value is a yaml file, we'll try to load it as yaml
-                    if value.endswith(".yaml") or value.endswith(".yml"):
+                    if value.endswith((".yaml", "yml")):
                         with open(value, "r") as f:
                             return yaml.safe_load(f)
                     with open(value, "r") as f:
@@ -398,8 +346,6 @@ class JsonParamType(click.ParamType):
     def convert(
         self, value: typing.Any, param: typing.Optional[click.Parameter], ctx: typing.Optional[click.Context]
     ) -> typing.Any:
-        if isinstance(value, ArtifactQuery):
-            return value
         if value is None:
             raise click.BadParameter("None value cannot be converted to a Json type.")
 
@@ -407,68 +353,38 @@ class JsonParamType(click.ParamType):
 
         # We compare the origin type because the json parsed value for list or dict is always a list or dict without
         # the covariant type information.
-        if type(parsed_value) == typing.get_origin(self._python_type) or type(parsed_value) == self._python_type:
+        if type(parsed_value) is typing.get_origin(self._python_type) or type(parsed_value) is self._python_type:
             # Indexing the return value of get_args will raise an error for native dict and list types.
             # We don't support native list/dict types with nested dataclasses.
             if get_args(self._python_type) == ():
                 return parsed_value
             elif isinstance(parsed_value, list) and dataclasses.is_dataclass(get_args(self._python_type)[0]):
                 j = JsonParamType(get_args(self._python_type)[0])
-                return [j.convert(v, param, ctx) for v in parsed_value]
+                # turn object back into json string
+                return [j.convert(json.dumps(v), param, ctx) for v in parsed_value]
             elif isinstance(parsed_value, dict) and dataclasses.is_dataclass(get_args(self._python_type)[1]):
                 j = JsonParamType(get_args(self._python_type)[1])
-                return {k: j.convert(v, param, ctx) for k, v in parsed_value.items()}
+                # turn object back into json string
+                return {k: j.convert(json.dumps(v), param, ctx) for k, v in parsed_value.items()}
 
             return parsed_value
 
-        if is_pydantic_basemodel(self._python_type):
-            """
-            This function supports backward compatibility for the Pydantic v1 plugin.
-            If the class is a Pydantic BaseModel, it attempts to parse JSON input using
-            the appropriate version of Pydantic (v1 or v2).
-            """
-            try:
-                if importlib.util.find_spec("pydantic.v1") is not None:
-                    from pydantic import BaseModel as BaseModelV2
-
-                    if issubclass(self._python_type, BaseModelV2):
-                        return self._python_type.model_validate_json(
-                            json.dumps(parsed_value), strict=False, context={"deserialize": True}
-                        )
-            except ImportError:
-                pass
-
-            # The behavior of the Pydantic v1 plugin.
-            return self._python_type.parse_raw(json.dumps(parsed_value))
-        return None
+        from pydantic import BaseModel
 
-
-def modify_literal_uris(lit: Literal):
-    """
-    Modifies the literal object recursively to replace the URIs with the native paths.
-    """
-    if lit.collection:
-        for l in lit.collection.literals:
-            modify_literal_uris(l)
-    elif lit.map:
-        for k, v in lit.map.literals.items():
-            modify_literal_uris(v)
-    elif lit.scalar:
-        if lit.scalar.blob and lit.scalar.blob.uri and lit.scalar.blob.uri.startswith(RemoteFSPathResolver.protocol):
-            lit.scalar.blob._uri = RemoteFSPathResolver.resolve_remote_path(lit.scalar.blob.uri)
-        elif lit.scalar.union:
-            modify_literal_uris(lit.scalar.union.value)
-        elif (
-            lit.scalar.structured_dataset
-            and lit.scalar.structured_dataset.uri
-            and lit.scalar.structured_dataset.uri.startswith(RemoteFSPathResolver.protocol)
-        ):
-            lit.scalar.structured_dataset._uri = RemoteFSPathResolver.resolve_remote_path(
-                lit.scalar.structured_dataset.uri
+        if issubclass(self._python_type, BaseModel):
+            return typing.cast(BaseModel, self._python_type).model_validate_json(
+                json.dumps(parsed_value), strict=False, context={"deserialize": True}
             )
+        elif dataclasses.is_dataclass(self._python_type):
+            from mashumaro.codecs.json import JSONDecoder
 
+            decoder = JSONDecoder(self._python_type)
+            return decoder.decode(value)
 
-SIMPLE_TYPE_CONVERTER: typing.Dict[SimpleType, click.ParamType] = {
+        return parsed_value
+
+
+SIMPLE_TYPE_CONVERTER = {
     SimpleType.FLOAT: click.FLOAT,
     SimpleType.INTEGER: click.INT,
     SimpleType.STRING: click.STRING,
@@ -482,7 +398,7 @@ def literal_type_to_click_type(lt: LiteralType, python_type: typing.Type) -> cli
     """
     Converts a Flyte LiteralType given a python_type to a click.ParamType
     """
-    if lt.simple:
+    if lt.HasField("simple"):
         if lt.simple == SimpleType.STRUCT:
             ct = JsonParamType(python_type)
             ct.name = f"JSON object {python_type.__name__}"
@@ -491,38 +407,38 @@ def literal_type_to_click_type(lt: LiteralType, python_type: typing.Type) -> cli
             return SIMPLE_TYPE_CONVERTER[lt.simple]
         raise NotImplementedError(f"Type {lt.simple} is not supported in pyflyte run")
 
-    if lt.enum_type:
-        return EnumParamType(python_type)  # type: ignore
-
-    if lt.structured_dataset_type:
+    if lt.HasField("structured_dataset_type"):
         return StructuredDatasetParamType()
 
-    if lt.collection_type or lt.map_value_type:
+    if lt.HasField("collection_type") or lt.HasField("map_value_type"):
         ct = JsonParamType(python_type)
-        if lt.collection_type:
+        if lt.HasField("collection_type"):
             ct.name = "json list"
         else:
             ct.name = "json dictionary"
         return ct
 
-    if lt.blob:
+    if lt.HasField("blob"):
         if lt.blob.dimensionality == BlobType.BlobDimensionality.SINGLE:
             if lt.blob.format == FlytePickleTransformer.PYTHON_PICKLE_FORMAT:
                 return PickleParamType()
+            # TODO: Add JSONIteratorTransformer
             # elif lt.blob.format == JSONIteratorTransformer.JSON_ITERATOR_FORMAT:
             #     return JSONIteratorParamType()
             return FileParamType()
         return DirParamType()
 
-    if lt.union_type:
+    if lt.HasField("union_type"):
         cts = []
         for i in range(len(lt.union_type.variants)):
             variant = lt.union_type.variants[i]
             variant_python_type = typing.get_args(python_type)[i]
-            ct = literal_type_to_click_type(variant, variant_python_type)
-            cts.append(ct)
+            cts.append(literal_type_to_click_type(variant, variant_python_type))
         return UnionParamType(cts)
 
+    if lt.HasField("enum_type"):
+        return EnumParamType(python_type)  # type: ignore
+
     return click.UNPROCESSED
 
 
@@ -533,9 +449,7 @@ class FlyteLiteralConverter(object):
         self,
         literal_type: LiteralType,
         python_type: typing.Type,
-        is_remote: bool,
     ):
-        self._is_remote = is_remote
         self._literal_type = literal_type
         self._python_type = python_type
         self._click_type = literal_type_to_click_type(literal_type, python_type)
@@ -551,25 +465,15 @@ class FlyteLiteralConverter(object):
         self, ctx: click.Context, param: typing.Optional[click.Parameter], value: typing.Any
     ) -> typing.Union[Literal, typing.Any]:
         """
-        Convert the value to a Flyte Literal or a python native type. This is used by click to convert the input.
+        Convert the value to a python native type. This is used by click to convert the input.
         """
-        if isinstance(value, ArtifactQuery):
-            return value
         try:
             # If the expected Python type is datetime.date, adjust the value to date
             if self._python_type is datetime.date:
                 # Click produces datetime, so converting to date to avoid type mismatch error
                 value = value.date()
-            # If the input matches the default value in the launch plan, serialization can be skipped.
-            if param and value == param.default:
-                return None
 
-            # If this is used for remote execution, then we need to convert it back to a python native type
-            if not self._is_remote:
-                return value
-
-            lit = asyncio.run(TypeEngine.to_literal(value, self._python_type, self._literal_type))
-            return lit
+            return value
         except click.BadParameter:
             raise
         except Exception as e:
@@ -577,3 +481,58 @@ class FlyteLiteralConverter(object):
                 f"Failed to convert param: {param if param else 'NA'}, value: {value} to type: {self._python_type}."
                 f" Reason {e}"
             ) from e
+
+
+def to_click_option(
+    input_name: str,
+    literal_var: Variable,
+    python_type: typing.Type,
+    default_val: typing.Any,
+) -> click.Option:
+    """
+    This handles converting workflow input types to supported click parameters with callbacks to initialize
+    the input values to their expected types.
+    """
+    from flyteidl.core.types_pb2 import SimpleType
+
+    if input_name != input_name.lower():
+        # Click does not support uppercase option names: https://github.com/pallets/click/issues/837
+        raise ValueError(f"Workflow input name must be lowercase: {input_name!r}")
+
+    literal_converter = FlyteLiteralConverter(
+        literal_type=literal_var.type,
+        python_type=python_type,
+    )
+
+    if literal_converter.is_bool() and not default_val:
+        default_val = False
+
+    description_extra = ""
+    if literal_var.type.simple == SimpleType.STRUCT:
+        if default_val:
+            # pydantic v2
+            if hasattr(default_val, "model_dump_json"):
+                default_val = default_val.model_dump_json()
+            else:
+                encoder = JSONEncoder(python_type)
+                default_val = encoder.encode(default_val)
+        if literal_var.type.metadata:
+            description_extra = f": {MessageToDict(literal_var.type.metadata)}"
+
+    # If a query has been specified, the input is never strictly required at this layer
+    required = False if default_val is not None else True
+    is_flag: typing.Optional[bool] = None
+    if literal_converter.is_bool():
+        required = False
+        is_flag = True
+
+    return click.Option(
+        param_decls=[f"--{input_name}"],
+        type=literal_converter.click_type,
+        is_flag=is_flag,
+        default=default_val,
+        show_default=True,
+        required=required,
+        help=literal_var.description + description_extra,
+        callback=literal_converter.convert,
+    )

flyte/{_cli → cli}/_run.py

@@ -1,22 +1,23 @@
 from __future__ import annotations
 
+import asyncio
+import inspect
 from dataclasses import dataclass, field, fields
 from pathlib import Path
 from types import ModuleType
 from typing import Any, Dict, List, cast
 
 import click
-from click import Context
+from click import Context, Parameter
 from rich.console import Console
 from typing_extensions import get_args
 
-import flyte
-
 from .._code_bundle._utils import CopyFiles
 from .._task import TaskTemplate
 from ..remote import Run
 from . import _common as common
 from ._common import CLIConfig
+from ._params import to_click_option
 
 
 @dataclass
@@ -48,6 +49,28 @@ class RunArguments:
             )
         },
     )
+    name: str | None = field(
+        default=None,
+        metadata={
+            "click.option": click.Option(
+                ["--name"],
+                type=str,
+                help="Name of the run. If not provided, a random name will be generated.",
+            )
+        },
+    )
+    follow: bool = field(
+        default=True,
+        metadata={
+            "click.option": click.Option(
+                ["--follow", "-f"],
+                is_flag=True,
+                default=False,
+                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.",
+            )
+        },
+    )
 
     @classmethod
     def from_dict(cls, d: Dict[str, Any]) -> RunArguments:
@@ -70,24 +93,60 @@ class RunTaskCommand(click.Command):
         super().__init__(obj_name, *args, **kwargs)
 
     def invoke(self, ctx: Context):
-        obj: CLIConfig = ctx.obj or CLIConfig()
+        obj: CLIConfig = ctx.obj
+        if obj is None:
+            import flyte.config
+
+            obj = CLIConfig(flyte.config.auto(), ctx)
+
         obj.init(self.run_args.project, self.run_args.domain)
 
-        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",
-        ).run(self.obj)
-        if isinstance(r, Run) and r.action is not None:
-            console = Console()
-            console.print(
-                common.get_panel(
-                    "Run",
-                    f"[green bold]Created Run: {r.name} [/green bold] "
-                    f"(Project: {r.action.action_id.run.project}, Domain: {r.action.action_id.run.domain})\n\n"
-                    f"[blue bold]{r.url}[/blue bold]",
+        async def _run():
+            import flyte
+
+            r = flyte.with_runcontext(
+                copy_style=self.run_args.copy_style,
+                mode="local" if self.run_args.local else "remote",
+                name=self.run_args.name,
+            ).run(self.obj, **ctx.params)
+            if isinstance(r, Run) and r.action is not None:
+                console = Console()
+                console.print(
+                    common.get_panel(
+                        "Run",
+                        f"[green bold]Created Run: {r.name} [/green bold] "
+                        f"(Project: {r.action.action_id.run.project}, Domain: {r.action.action_id.run.domain})\n"
+                        f"➡️  [blue bold]{r.url}[/blue bold]",
+                    )
                 )
-            )
+                if self.run_args.follow:
+                    console.print(
+                        "[dim]Log streaming enabled, will wait for task to start running "
+                        "and log stream to be available[/dim]"
+                    )
+                    await r.show_logs(max_lines=30, show_ts=True, raw=False)
+
+        asyncio.run(_run())
+
+    def get_params(self, ctx: Context) -> List[Parameter]:
+        # Note this function may be called multiple times by click.
+        task = self.obj
+        from .._internal.runtime.types_serde import transform_native_to_typed_interface
+
+        interface = transform_native_to_typed_interface(task.native_interface)
+        if interface is None:
+            return super().get_params(ctx)
+        inputs_interface = task.native_interface.inputs
+
+        params: List[Parameter] = []
+        for name, var in interface.inputs.variables.items():
+            default_val = None
+            if inputs_interface[name][1] is not inspect._empty:
+                default_val = inputs_interface[name][1]
+            params.append(to_click_option(name, var, inputs_interface[name][0], default_val))
+
+        self.params = params
+        return super().get_params(ctx)
 
 
 class TaskPerFileGroup(common.ObjectsPerFileGroup):
@@ -142,11 +201,31 @@ class TaskFiles(common.FileGroup):
             filename=Path(filename),
             run_args=run_args,
             name=filename,
-            help=f"Run, functions decorated `env.task` or instances of Tasks in {filename}",
+            help=f"Run, functions decorated with `env.task` in {filename}",
         )
 
 
 run = TaskFiles(
     name="run",
-    help="Run a task from a python file",
+    help="""
+Run a task from a python file.
+
+Example usage:
+```bash
+flyte run --name examples/basics/hello.py my_task --arg1 value1 --arg2 value2
+```
+Note: all arguments for the run command are provided right after the `run` command and before the file name.
+
+You can also specify the project and domain using the `--project` and `--domain` options, respectively. These
+options can be set in the config file or passed as command line arguments.
+
+Note: The arguments for the task are provided after the task name and can be retrieved using `--help`
+Example:
+```bash
+flyte run --name examples/basics/hello.py my_task --help
+```
+
+To run a task locally, use the `--local` flag. This will run the task in the local environment instead of the remote
+ Flyte environment.
+""",
 )