flyte 0.0.1b1__py3-none-any.whl → 0.0.1b2__py3-none-any.whl

union/_cli/_get.py

@@ -0,0 +1,162 @@
+import asyncio
+
+import rich_click as click
+from rich.console import Console
+
+from . import _common as common
+
+
+@click.group(name="get")
+def get():
+    """
+    Get the value of a task or environment.
+    """
+
+
+@get.command()
+@click.argument("name", type=str, required=False)
+@click.pass_obj
+def project(cfg: common.CLIConfig, name: str | None = None):
+    """
+    Get the current project.
+    """
+    from union.remote import Project
+
+    print(cfg)
+    cfg.init()
+
+    console = Console()
+    if name:
+        console.print(Project.get(name))
+    else:
+        console.print(common.get_table("Projects", Project.listall()))
+
+
+@get.command(cls=common.CommandBase)
+@click.argument("name", type=str, required=False)
+@click.pass_obj
+def run(cfg: common.CLIConfig, name: str | None = None, project: str | None = None, domain: str | None = None):
+    """
+    Get the current run.
+    """
+    from union.remote import Run, RunDetails
+
+    cfg.init(project=project, domain=domain)
+
+    console = Console()
+    if name:
+        details = RunDetails.get(name=name)
+        console.print(details)
+    else:
+        console.print(common.get_table("Runs", Run.listall()))
+
+
+@get.command(cls=common.CommandBase)
+@click.argument("name", type=str, required=False)
+@click.argument("version", type=str, required=False)
+@click.pass_obj
+def task(
+    cfg: common.CLIConfig,
+    name: str | None = None,
+    version: str | None = None,
+    project: str | None = None,
+    domain: str | None = None,
+):
+    """
+    Get the current task.
+    """
+    from union.remote import Task
+
+    cfg.init(project=project, domain=domain)
+
+    console = Console()
+    if name:
+        v = Task.get(name=name, version=version)
+        if v is None:
+            raise click.BadParameter(f"Task {name} not found.")
+        t = v.fetch()
+        console.print(t)
+    else:
+        raise click.BadParameter("Task listing is not supported yet, please provide a name.")
+        # console.print(common.get_table("Tasks", Task.listall()))
+
+
+@get.command(cls=common.CommandBase)
+@click.argument("run_name", type=str, required=True)
+@click.argument("action_name", type=str, required=False)
+@click.pass_obj
+def action(
+    cfg: common.CLIConfig,
+    run_name: str,
+    action_name: str | None = None,
+    project: str | None = None,
+    domain: str | None = None,
+):
+    """
+    Get all actions for a run or details for a specific action.
+    """
+    import union.remote as remote
+
+    cfg.init(project=project, domain=domain)
+
+    console = Console()
+    if action_name:
+        console.print(remote.Action.get(run_name=run_name, name=action_name))
+    else:
+        # List all actions for the run
+        console.print(common.get_table(f"Actions for {run_name}", remote.Action.listall(for_run_name=run_name)))
+
+
+@get.command(cls=common.CommandBase)
+@click.argument("run_name", type=str, required=False)
+@click.argument("action_name", type=str, required=False)
+@click.pass_obj
+def logs(
+    cfg: common.CLIConfig,
+    run_name: str,
+    action_name: str | None = None,
+    project: str | None = None,
+    domain: str | None = None,
+):
+    """
+    Get the current run.
+    """
+    import union.remote as remote
+
+    cfg.init(project=project, domain=domain)
+
+    async def _run_log_view(_obj):
+        task = asyncio.create_task(_obj.show_logs())
+        try:
+            await task
+        except KeyboardInterrupt:
+            task.cancel()
+
+    if action_name:
+        obj = remote.Action.get(run_name=run_name, name=action_name)
+    else:
+        obj = remote.Run.get(run_name)
+    asyncio.run(_run_log_view(obj))
+
+
+@get.command(cls=common.CommandBase)
+@click.argument("name", type=str, required=False)
+@click.pass_obj
+def secret(
+    cfg: common.CLIConfig,
+    name: str | None = None,
+    project: str | None = None,
+    domain: str | None = None,
+):
+    """
+    Get the current secret.
+    """
+    import union.remote as remote
+
+    cfg.init(project=project, domain=domain)
+
+    console = Console()
+    if name:
+        console.print(remote.Secret.get(name))
+    else:
+        console.print(common.get_table("Secrets", remote.Secret.listall()))

{flyte → union}/_cli/_params.py

@@ -1,3 +1,4 @@
+import asyncio
 import dataclasses
 import datetime
 import enum
@@ -6,7 +7,6 @@ import importlib.util
 import json
 import os
 import pathlib
-import re
 import sys
 import typing
 import typing as t
@@ -14,16 +14,24 @@ 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 flyte._logging import logger
-from flyte.io import Dir, File
-from flyte.io.pickle.transformer import FlytePickleTransformer
+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
 
 
 class StructuredDataset:
@@ -35,6 +43,26 @@ 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.
@@ -72,13 +100,17 @@ class DirParamType(click.ParamType):
     def convert(
         self, value: typing.Any, param: typing.Optional[click.Parameter], ctx: typing.Optional[click.Context]
     ) -> typing.Any:
-        from flyte.storage import is_remote
+        if isinstance(value, ArtifactQuery):
+            return value
 
+        # 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)
+        return Dir(path=value, remote_directory=remote_directory)
 
 
 class StructuredDatasetParamType(click.ParamType):
@@ -91,6 +123,8 @@ 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):
@@ -104,19 +138,22 @@ class FileParamType(click.ParamType):
     def convert(
         self, value: typing.Any, param: typing.Optional[click.Parameter], ctx: typing.Optional[click.Context]
     ) -> typing.Any:
-        from flyte.storage import is_remote
-
+        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
         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)
+        return File(path=value, remote_path=remote_path)
 
 
 class PickleParamType(click.ParamType):
     name = "pickle"
 
-    def get_metavar(self, param: "Parameter", *args) -> t.Optional[str]:
+    def get_metavar(self, param: click.Parameter) -> t.Optional[str]:
         return "Python Object <Module>:<Object>"
 
     def convert(
@@ -126,7 +163,7 @@ class PickleParamType(click.ParamType):
             return value
         parts = value.split(":")
         if len(parts) != 2:
-            if ctx and ctx.obj and ctx.obj.log_level >= 10:  # DEBUG level
+            if ctx and ctx.obj and ctx.obj.verbose > 0:
                 click.echo(f"Did not receive a string in the expected format <MODULE>:<VAR>, falling back to: {value}")
             return value
         try:
@@ -148,6 +185,9 @@ 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'
@@ -171,10 +211,10 @@ def parse_human_durations(text: str) -> list[datetime.timedelta]:
     durations = []
 
     for part in raw_parts:
-        new_part = part.strip().lower()
+        part = part.strip().lower()
 
         # Match 1:24 or :45
-        m_colon = re.match(r"^(?:(\d+):)?(\d+)$", new_part)
+        m_colon = re.match(r"^(?:(\d+):)?(\d+)$", part)
         if m_colon:
             minutes = int(m_colon.group(1)) if m_colon.group(1) else 0
             seconds = int(m_colon.group(2))
@@ -182,7 +222,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?$", new_part)
+        m_units = re.match(r"^(\d+)\s*(day|hour|minute|second)s?$", part)
         if m_units:
             value = int(m_units.group(1))
             unit = m_units.group(2)
@@ -230,6 +270,9 @@ 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
 
@@ -260,6 +303,8 @@ 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)
@@ -273,6 +318,8 @@ 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))
@@ -286,7 +333,10 @@ class UnionParamType(click.ParamType):
     def __init__(self, types: typing.List[click.ParamType]):
         super().__init__()
         self._types = self._sort_precedence(types)
-        self.name = "|".join([t.name for t in self._types])
+
+    @property
+    def name(self) -> str:
+        return "|".join([t.name for t in self._types])
 
     @staticmethod
     def _sort_precedence(tp: typing.List[click.ParamType]) -> typing.List[click.ParamType]:
@@ -300,7 +350,7 @@ class UnionParamType(click.ParamType):
                 str_types.append(p)
             else:
                 others.append(p)
-        return others + str_types + unprocessed  # type: ignore
+        return others + str_types + unprocessed
 
     def convert(
         self, value: typing.Any, param: typing.Optional[click.Parameter], ctx: typing.Optional[click.Context]
@@ -309,6 +359,8 @@ 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)
@@ -334,7 +386,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", "yml")):
+                    if value.endswith(".yaml") or value.endswith(".yml"):
                         with open(value, "r") as f:
                             return yaml.safe_load(f)
                     with open(value, "r") as f:
@@ -346,6 +398,8 @@ 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.")
 
@@ -353,38 +407,68 @@ 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) is typing.get_origin(self._python_type) or type(parsed_value) is self._python_type:
+        if type(parsed_value) == typing.get_origin(self._python_type) or type(parsed_value) == 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])
-                # turn object back into json string
-                return [j.convert(json.dumps(v), param, ctx) for v in parsed_value]
+                return [j.convert(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])
-                # turn object back into json string
-                return {k: j.convert(json.dumps(v), param, ctx) for k, v in parsed_value.items()}
+                return {k: j.convert(v, param, ctx) for k, v in parsed_value.items()}
 
             return parsed_value
 
-        from pydantic import BaseModel
-
-        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
+        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
 
-            decoder = JSONDecoder(self._python_type)
-            return decoder.decode(value)
 
-        return parsed_value
+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
+            )
 
 
-SIMPLE_TYPE_CONVERTER = {
+SIMPLE_TYPE_CONVERTER: typing.Dict[SimpleType, click.ParamType] = {
     SimpleType.FLOAT: click.FLOAT,
     SimpleType.INTEGER: click.INT,
     SimpleType.STRING: click.STRING,
@@ -398,7 +482,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.HasField("simple"):
+    if lt.simple:
         if lt.simple == SimpleType.STRUCT:
             ct = JsonParamType(python_type)
             ct.name = f"JSON object {python_type.__name__}"
@@ -407,38 +491,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.HasField("structured_dataset_type"):
+    if lt.enum_type:
+        return EnumParamType(python_type)  # type: ignore
+
+    if lt.structured_dataset_type:
         return StructuredDatasetParamType()
 
-    if lt.HasField("collection_type") or lt.HasField("map_value_type"):
+    if lt.collection_type or lt.map_value_type:
         ct = JsonParamType(python_type)
-        if lt.HasField("collection_type"):
+        if lt.collection_type:
             ct.name = "json list"
         else:
             ct.name = "json dictionary"
         return ct
 
-    if lt.HasField("blob"):
+    if lt.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.HasField("union_type"):
+    if lt.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]
-            cts.append(literal_type_to_click_type(variant, variant_python_type))
+            ct = literal_type_to_click_type(variant, variant_python_type)
+            cts.append(ct)
         return UnionParamType(cts)
 
-    if lt.HasField("enum_type"):
-        return EnumParamType(python_type)  # type: ignore
-
     return click.UNPROCESSED
 
 
@@ -449,7 +533,9 @@ 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)
@@ -465,15 +551,25 @@ 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 python native type. This is used by click to convert the input.
+        Convert the value to a Flyte Literal or 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
 
-            return value
+            # 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
         except click.BadParameter:
             raise
         except Exception as e:
@@ -481,58 +577,3 @@ 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,
-    )