flyte 2.0.0b13__py3-none-any.whl → 2.0.0b30__py3-none-any.whl

flyte/cli/_serve.py

@@ -0,0 +1,64 @@
+from typing import List
+
+import click
+
+
+@click.group("serve")
+@click.pass_context
+def serve(_: click.Context):
+    """
+    Start the specific service. For example:
+
+    ```bash
+    flyte serve connector
+    ```
+    """
+
+
+@serve.command()
+@click.option(
+    "--port",
+    default="8000",
+    is_flag=False,
+    type=int,
+    help="Grpc port for the connector service",
+)
+@click.option(
+    "--prometheus_port",
+    default="9090",
+    is_flag=False,
+    type=int,
+    help="Prometheus port for the connector service",
+)
+@click.option(
+    "--worker",
+    default="10",
+    is_flag=False,
+    type=int,
+    help="Number of workers for the grpc server",
+)
+@click.option(
+    "--timeout",
+    default=None,
+    is_flag=False,
+    type=int,
+    help="It will wait for the specified number of seconds before shutting down grpc server. It should only be used "
+    "for testing.",
+)
+@click.option(
+    "--modules",
+    required=False,
+    multiple=True,
+    type=str,
+    help="List of additional files or module that defines the connector",
+)
+@click.pass_context
+def connector(
+    _: click.Context, port: int, prometheus_port: int, worker: int, timeout: int | None, modules: List[str] | None
+):
+    """
+    Start a grpc server for the connector service.
+    """
+    from flyte.connectors import ConnectorService
+
+    ConnectorService.run(port, prometheus_port, worker, timeout, modules)

flyte/cli/_update.py

@@ -0,0 +1,37 @@
+import rich_click as click
+
+import flyte.remote as remote
+
+from . import _common as common
+
+
+@click.group(name="update")
+def update():
+    """
+    Update various flyte entities.
+    """
+
+
+@update.command("trigger", cls=common.CommandBase)
+@click.argument("name", type=str)
+@click.argument("task_name", type=str)
+@click.option("--activate/--deactivate", required=True, help="Activate or deactivate the trigger.")
+@click.pass_obj
+def trigger(cfg: common.CLIConfig, name: str, task_name: str, activate: bool, project: str | None, domain: str | None):
+    """
+    Update a trigger.
+
+    \b
+    Example usage:
+
+    ```bash
+    flyte update trigger <trigger_name> <task_name> --activate | --deactivate
+    [--project <project_name> --domain <domain_name>]
+    ```
+    """
+    cfg.init(project, domain)
+    console = common.get_console()
+    to_state = "active" if activate else "deactivate"
+    with console.status(f"Updating trigger {name} for task {task_name} to {to_state}..."):
+        remote.Trigger.update(name, task_name, activate)
+    console.print(f"Trigger updated and is set to [fuchsia]{to_state}[/fuchsia]")

flyte/cli/_user.py

@@ -0,0 +1,17 @@
+import rich_click as click
+
+import flyte.remote as remote
+
+from . import _common as common
+
+
+@click.command()
+@click.pass_obj
+def whoami(
+    cfg: common.CLIConfig,
+):
+    """Display the current user information."""
+    cfg.init()
+    console = common.get_console()
+    user_info = remote.User.get()
+    console.print(user_info.to_json())

flyte/cli/main.py

@@ -1,7 +1,7 @@
 import rich_click as click
 from typing_extensions import get_args
 
-from flyte._logging import initialize_logger, logger
+from flyte._logging import LogFormat, initialize_logger, logger
 
 from . import _common as common
 from ._abort import abort
@@ -12,7 +12,11 @@ from ._delete import delete
 from ._deploy import deploy
 from ._gen import gen
 from ._get import get
+from ._plugins import discover_and_register_plugins
 from ._run import run
+from ._serve import serve
+from ._update import update
+from ._user import whoami
 
 help_config = click.RichHelpConfiguration(
     use_markdown=True,
@@ -25,7 +29,7 @@ help_config = click.RichHelpConfiguration(
             },
             {
                 "name": "Management of various objects.",
-                "commands": ["create", "get", "delete"],
+                "commands": ["create", "get", "delete", "update"],
             },
             {
                 "name": "Build and deploy environments, tasks and images.",
@@ -35,6 +39,10 @@ help_config = click.RichHelpConfiguration(
                 "name": "Documentation generation",
                 "commands": ["gen"],
             },
+            {
+                "name": "User information",
+                "commands": ["whoami"],
+            },
         ]
     },
 )
@@ -116,6 +124,16 @@ def _verbosity_to_loglevel(verbosity: int) -> int | None:
     show_default=True,
     required=False,
 )
+@click.option(
+    "--log-format",
+    type=click.Choice(get_args(LogFormat), case_sensitive=False),
+    envvar="LOG_FORMAT",
+    default="console",
+    help="Formatting for logs, defaults to 'console' which is meant to be human readable."
+    " 'json' is meant for machine parsing.",
+    show_default=True,
+    required=False,
+)
 @click.rich_config(help_config=help_config)
 @click.pass_context
 def main(
@@ -123,6 +141,7 @@ def main(
     endpoint: str | None,
     insecure: bool,
     verbose: int,
+    log_format: LogFormat,
     org: str | None,
     config_file: str | None,
     auth_type: str | None = None,
@@ -166,8 +185,8 @@ def main(
     import flyte.config as config
 
     log_level = _verbosity_to_loglevel(verbose)
-    if log_level is not None:
-        initialize_logger(log_level)
+    if log_level is not None or log_format != "console":
+        initialize_logger(log_level=log_level, log_format=log_format)
 
     cfg = config.auto(config_file=config_file)
     if cfg.source:
@@ -175,6 +194,7 @@ def main(
 
     ctx.obj = CLIConfig(
         log_level=log_level,
+        log_format=log_format,
         endpoint=endpoint,
         insecure=insecure,
         org=org,
@@ -193,3 +213,9 @@ main.add_command(abort)  # type: ignore
 main.add_command(gen)  # type: ignore
 main.add_command(delete)  # type: ignore
 main.add_command(build)
+main.add_command(whoami)  # type: ignore
+main.add_command(update)  # type: ignore
+main.add_command(serve)  # type: ignore
+
+# Discover and register CLI plugins from installed packages
+discover_and_register_plugins(main)

flyte/config/_config.py

@@ -148,6 +148,7 @@ class ImageConfig(object):
     """
 
     builder: str | None = None
+    image_refs: typing.Dict[str, str] = field(default_factory=dict)
 
     @classmethod
     def auto(cls, config_file: typing.Optional[typing.Union[str, ConfigFile]] = None) -> "ImageConfig":
@@ -159,6 +160,7 @@ class ImageConfig(object):
         config_file = get_config_file(config_file)
         kwargs: typing.Dict[str, typing.Any] = {}
         kwargs = set_if_exists(kwargs, "builder", _internal.Image.BUILDER.read(config_file))
+        kwargs = set_if_exists(kwargs, "image_refs", _internal.Image.IMAGE_REFS.read(config_file))
         return ImageConfig(**kwargs)
 
 
@@ -192,7 +194,7 @@ class Config(object):
         )
 
     @classmethod
-    def auto(cls, config_file: typing.Union[str, ConfigFile, None] = None) -> "Config":
+    def auto(cls, config_file: typing.Union[str, pathlib.Path, ConfigFile, None] = None) -> "Config":
         """
         Automatically constructs the Config Object. The order of precedence is as follows
           1. first try to find any env vars that match the config vars specified in the FLYTE_CONFIG format.
@@ -225,16 +227,18 @@ def set_if_exists(d: dict, k: str, val: typing.Any) -> dict:
     return d
 
 
-def auto(config_file: typing.Union[str, ConfigFile, None] = None) -> Config:
+def auto(config_file: typing.Union[str, pathlib.Path, ConfigFile, None] = None) -> Config:
     """
     Automatically constructs the Config Object. The order of precedence is as follows
       1. If specified, read the config from the provided file path.
       2. If not specified, the config file is searched in the default locations.
             a. ./config.yaml if it exists  (current working directory)
-            b. `UCTL_CONFIG` environment variable
-            c. `FLYTECTL_CONFIG` environment variable
-            d. ~/.union/config.yaml if it exists
-            e. ~/.flyte/config.yaml if it exists
+            b. ./.flyte/config.yaml if it exists (current working directory)
+            c. <git_root>/.flyte/config.yaml if it exists
+            d. `UCTL_CONFIG` environment variable
+            e. `FLYTECTL_CONFIG` environment variable
+            f. ~/.union/config.yaml if it exists
+            g. ~/.flyte/config.yaml if it exists
     3. If any value is not found in the config file, the default value is used.
     4. For any value there are environment variables that match the config variable names, those will override
 

flyte/config/_internal.py

@@ -70,3 +70,4 @@ class Image(object):
     """
 
     BUILDER = ConfigEntry(YamlConfigEntry("image.builder"))
+    IMAGE_REFS = ConfigEntry(YamlConfigEntry("image.image_refs"))

flyte/config/_reader.py

@@ -108,7 +108,7 @@ class ConfigFile(object):
         return pathlib.Path(self._location)
 
     @staticmethod
-    def _read_yaml_config(location: str) -> typing.Optional[typing.Dict[str, typing.Any]]:
+    def _read_yaml_config(location: str | pathlib.Path) -> typing.Optional[typing.Dict[str, typing.Any]]:
         with open(location, "r") as fh:
             try:
                 yaml_contents = yaml.safe_load(fh)
@@ -135,20 +135,41 @@ class ConfigFile(object):
         return self._yaml_config
 
 
+def _config_path_from_git_root() -> pathlib.Path | None:
+    from flyte.git import config_from_root
+
+    config = config_from_root()
+    if config is None:
+        return None
+    return config.source
+
+
 def resolve_config_path() -> pathlib.Path | None:
     """
     Config is read from the following locations in order of precedence:
     1. ./config.yaml if it exists
-    2. `UCTL_CONFIG` environment variable
-    3. `FLYTECTL_CONFIG` environment variable
-    4. ~/.union/config.yaml if it exists
-    5. ~/.flyte/config.yaml if it exists
+    2. ./.flyte/config.yaml if it exists
+    3. <git_root>/.flyte/config.yaml if it exists
+    4. `UCTL_CONFIG` environment variable
+    5. `FLYTECTL_CONFIG` environment variable
+    6. ~/.union/config.yaml if it exists
+    7. ~/.flyte/config.yaml if it exists
     """
     current_location_config = Path("config.yaml")
     if current_location_config.exists():
         return current_location_config
     logger.debug("No ./config.yaml found")
 
+    dot_flyte_config = Path(".flyte", "config.yaml")
+    if dot_flyte_config.exists():
+        return dot_flyte_config
+    logger.debug("No ./.flyte/config.yaml found")
+
+    git_root_config = _config_path_from_git_root()
+    if git_root_config:
+        return git_root_config
+    logger.debug("No .flyte/config.yaml found in git repo root")
+
     uctl_path_from_env = getenv(UCTL_CONFIG_ENV_VAR, None)
     if uctl_path_from_env:
         return pathlib.Path(uctl_path_from_env)
@@ -173,13 +194,13 @@ def resolve_config_path() -> pathlib.Path | None:
 
 
 @lru_cache
-def get_config_file(c: typing.Union[str, ConfigFile, None]) -> ConfigFile | None:
+def get_config_file(c: typing.Union[str, pathlib.Path, ConfigFile, None]) -> ConfigFile | None:
     """
     Checks if the given argument is a file or a configFile and returns a loaded configFile else returns None
     """
-    if isinstance(c, str):
+    if isinstance(c, (str, pathlib.Path)):
         logger.debug(f"Using specified config file at {c}")
-        return ConfigFile(c)
+        return ConfigFile(str(c))
     elif isinstance(c, ConfigFile):
         return c
     config_path = resolve_config_path()

flyte/connectors/__init__.py

@@ -0,0 +1,11 @@
+from ._connector import AsyncConnector, AsyncConnectorExecutorMixin, ConnectorRegistry, Resource, ResourceMeta
+from ._server import ConnectorService
+
+__all__ = [
+    "AsyncConnector",
+    "AsyncConnectorExecutorMixin",
+    "ConnectorRegistry",
+    "ConnectorService",
+    "Resource",
+    "ResourceMeta",
+]

flyte/connectors/_connector.py

@@ -0,0 +1,270 @@
+import asyncio
+import json
+import typing
+from abc import ABC, abstractmethod
+from dataclasses import asdict, dataclass
+from typing import Any, Dict, List, Optional
+
+from flyteidl2.core import tasks_pb2
+from flyteidl2.core.execution_pb2 import TaskExecution, TaskLog
+from flyteidl2.plugins import connector_pb2
+from flyteidl2.plugins.connector_pb2 import Connector as ConnectorProto
+from flyteidl2.plugins.connector_pb2 import (
+    GetTaskLogsResponse,
+    GetTaskMetricsResponse,
+    TaskCategory,
+    TaskExecutionMetadata,
+)
+from google.protobuf import json_format
+from google.protobuf.struct_pb2 import Struct
+
+from flyte import Secret
+from flyte._context import internal_ctx
+from flyte._initialize import get_init_config
+from flyte._internal.runtime.convert import convert_from_native_to_outputs
+from flyte._internal.runtime.task_serde import get_proto_task
+from flyte._logging import logger
+from flyte._task import TaskTemplate
+from flyte.connectors.utils import is_terminal_phase
+from flyte.models import NativeInterface, SerializationContext
+from flyte.types._type_engine import dataclass_from_dict
+
+
+@dataclass(frozen=True)
+class ConnectorRegistryKey:
+    task_type_name: str
+    task_type_version: int
+
+
+@dataclass
+class ResourceMeta:
+    """
+    This is the metadata for the job. For example, the id of the job.
+    """
+
+    def encode(self) -> bytes:
+        """
+        Encode the resource meta to bytes.
+        """
+        return json.dumps(asdict(self)).encode("utf-8")
+
+    @classmethod
+    def decode(cls, data: bytes) -> "ResourceMeta":
+        """
+        Decode the resource meta from bytes.
+        """
+        return dataclass_from_dict(cls, json.loads(data.decode("utf-8")))
+
+
+@dataclass
+class Resource:
+    """
+    This is the output resource of the job.
+
+    Attributes
+    ----------
+        phase : TaskExecution.Phase
+            The phase of the job.
+        message : Optional[str]
+            The return message from the job.
+        log_links : Optional[List[TaskLog]]
+            The log links of the job. For example, the link to the BigQuery Console.
+        outputs : Optional[Union[LiteralMap, typing.Dict[str, Any]]]
+            The outputs of the job. If return python native types, the agent will convert them to flyte literals.
+        custom_info : Optional[typing.Dict[str, Any]]
+            The custom info of the job. For example, the job config.
+    """
+
+    phase: TaskExecution.Phase
+    message: Optional[str] = None
+    log_links: Optional[List[TaskLog]] = None
+    outputs: Optional[Dict[str, Any]] = None
+    custom_info: Optional[typing.Dict[str, Any]] = None
+
+
+class AsyncConnector(ABC):
+    """
+    This is the base class for all async connectors, and it defines the interface that all connectors must implement.
+    The connector service is responsible for invoking connectors.
+    The executor will communicate with the connector service to create tasks, get the status of tasks, and delete tasks.
+
+    All the connectors should be registered in the ConnectorRegistry.
+    Connector Service will look up the connector based on the task type and version.
+    """
+
+    name = "Async Connector"
+    task_type_name: str
+    task_type_version: int = 0
+    metadata_type: ResourceMeta
+
+    @abstractmethod
+    async def create(
+        self,
+        task_template: tasks_pb2.TaskTemplate,
+        output_prefix: str,
+        inputs: Optional[Dict[str, typing.Any]] = None,
+        task_execution_metadata: Optional[TaskExecutionMetadata] = None,
+        **kwargs,
+    ) -> ResourceMeta:
+        """
+        Return a resource meta that can be used to get the status of the task.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def get(self, resource_meta: ResourceMeta, **kwargs) -> Resource:
+        """
+        Return the status of the task, and return the outputs in some cases. For example, bigquery job
+        can't write the structured dataset to the output location, so it returns the output literals to the propeller,
+        and the propeller will write the structured dataset to the blob store.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def delete(self, resource_meta: ResourceMeta, **kwargs):
+        """
+        Delete the task. This call should be idempotent. It should raise an error if fails to delete the task.
+        """
+        raise NotImplementedError
+
+    async def get_metrics(self, resource_meta: ResourceMeta, **kwargs) -> GetTaskMetricsResponse:
+        """
+        Return the metrics for the task.
+        """
+        raise NotImplementedError
+
+    async def get_logs(self, resource_meta: ResourceMeta, **kwargs) -> GetTaskLogsResponse:
+        """
+        Return the metrics for the task.
+        """
+        raise NotImplementedError
+
+
+class ConnectorRegistry(object):
+    """
+    This is the registry for all connectors.
+    The connector service will look up the connector registry based on the task type and version.
+    """
+
+    _REGISTRY: typing.ClassVar[Dict[ConnectorRegistryKey, AsyncConnector]] = {}
+    _METADATA: typing.ClassVar[Dict[str, ConnectorProto]] = {}
+
+    @staticmethod
+    def register(connector: AsyncConnector, override: bool = False):
+        key = ConnectorRegistryKey(
+            task_type_name=connector.task_type_name, task_type_version=connector.task_type_version
+        )
+        if key in ConnectorRegistry._REGISTRY and override is False:
+            raise ValueError(
+                f"Duplicate connector for task type: {connector.task_type_name}"
+                f" and version: {connector.task_type_version}"
+            )
+        ConnectorRegistry._REGISTRY[key] = connector
+
+        task_category = TaskCategory(name=connector.task_type_name, version=connector.task_type_version)
+
+        if connector.name in ConnectorRegistry._METADATA:
+            connector_metadata = ConnectorRegistry._get_connector_metadata(connector.name)
+            connector_metadata.supported_task_categories.append(task_category)
+        else:
+            connector_metadata = ConnectorProto(
+                name=connector.name,
+                supported_task_categories=[task_category],
+            )
+            ConnectorRegistry._METADATA[connector.name] = connector_metadata
+
+    @staticmethod
+    def get_connector(task_type_name: str, task_type_version: int = 0) -> AsyncConnector:
+        key = ConnectorRegistryKey(task_type_name=task_type_name, task_type_version=task_type_version)
+        if key not in ConnectorRegistry._REGISTRY:
+            raise FlyteConnectorNotFound(
+                f"Cannot find connector for task type: {task_type_name} and version: {task_type_version}"
+            )
+        return ConnectorRegistry._REGISTRY[key]
+
+    @staticmethod
+    def _list_connectors() -> List[ConnectorProto]:
+        return list(ConnectorRegistry._METADATA.values())
+
+    @staticmethod
+    def _get_connector_metadata(name: str) -> ConnectorProto:
+        if name not in ConnectorRegistry._METADATA:
+            raise FlyteConnectorNotFound(f"Cannot find connector for name: {name}.")
+        return ConnectorRegistry._METADATA[name]
+
+
+class ConnectorSecretsMixin:
+    def __init__(self, secrets: Dict[str, str]):
+        # Key is the id of the secret, value is the secret name.
+        self._secrets = secrets
+
+    @property
+    def secrets(self) -> List[Secret]:
+        return [Secret(key=k, as_env_var=v) for k, v in self._secrets.items()]
+
+
+class AsyncConnectorExecutorMixin:
+    """
+    This mixin class is used to run the connector task locally, and it's only used for local execution.
+    Task should inherit from this class if the task can be run in the connector.
+    """
+
+    async def execute(self, **kwargs) -> Any:
+        task = typing.cast(TaskTemplate, self)
+        connector = ConnectorRegistry.get_connector(task.task_type, task.task_type_version)
+
+        ctx = internal_ctx()
+        tctx = internal_ctx().data.task_context
+        cfg = get_init_config()
+
+        if tctx is None:
+            raise RuntimeError("Task context is not set.")
+
+        sc = SerializationContext(
+            project=tctx.action.project,
+            domain=tctx.action.domain,
+            org=tctx.action.org,
+            code_bundle=tctx.code_bundle,
+            version=tctx.version,
+            image_cache=tctx.compiled_image_cache,
+            root_dir=cfg.root_dir,
+        )
+        tt = get_proto_task(task, sc)
+        resource_meta = await connector.create(task_template=tt, output_prefix=ctx.raw_data.path, inputs=kwargs)
+        resource = Resource(phase=TaskExecution.RUNNING)
+
+        while not is_terminal_phase(resource.phase):
+            resource = await connector.get(resource_meta=resource_meta)
+
+            if resource.log_links:
+                for link in resource.log_links:
+                    logger.info(f"{link.name}: {link.uri}")
+            await asyncio.sleep(1)
+
+        if resource.phase != TaskExecution.SUCCEEDED:
+            raise RuntimeError(f"Failed to run the task {task.name} with error: {resource.message}")
+
+        # TODO: Support abort
+
+        if resource.outputs is None:
+            return None
+        return tuple(resource.outputs.values())
+
+
+async def get_resource_proto(resource: Resource) -> connector_pb2.Resource:
+    if resource.outputs:
+        interface = NativeInterface.from_types(inputs={}, outputs={k: type(v) for k, v in resource.outputs.items()})
+        outputs = await convert_from_native_to_outputs(tuple(resource.outputs.values()), interface)
+    else:
+        outputs = None
+
+    return connector_pb2.Resource(
+        phase=resource.phase,
+        message=resource.message,
+        log_links=resource.log_links,
+        outputs=outputs,
+        custom_info=(json_format.Parse(json.dumps(resource.custom_info), Struct()) if resource.custom_info else None),
+    )
+
+
+class FlyteConnectorNotFound(ValueError): ...