flyte 0.2.0b23__py3-none-any.whl → 0.2.0b25__py3-none-any.whl

flyte/_run.py

@@ -6,13 +6,9 @@ import uuid
 from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union, cast
 
 import flyte.errors
-from flyte.errors import InitializationError
-from flyte.models import ActionID, Checkpoints, CodeBundle, RawDataPath, SerializationContext, TaskContext
-from flyte.syncify import syncify
-
-from ._context import contextual_run, internal_ctx
-from ._environment import Environment
-from ._initialize import (
+from flyte._context import contextual_run, internal_ctx
+from flyte._environment import Environment
+from flyte._initialize import (
     _get_init_config,
     get_client,
     get_common_config,
@@ -20,9 +16,19 @@ from ._initialize import (
     requires_initialization,
     requires_storage,
 )
-from ._logging import logger
-from ._task import P, R, TaskTemplate
-from ._tools import ipython_check
+from flyte._logging import logger
+from flyte._task import P, R, TaskTemplate
+from flyte._tools import ipython_check
+from flyte.errors import InitializationError
+from flyte.models import (
+    ActionID,
+    Checkpoints,
+    CodeBundle,
+    RawDataPath,
+    SerializationContext,
+    TaskContext,
+)
+from flyte.syncify import syncify
 
 if TYPE_CHECKING:
     from flyte.remote import Run
@@ -132,7 +138,9 @@ class _Runner:
 
             if self._interactive_mode:
                 code_bundle = await build_pkl_bundle(
-                    obj, upload_to_controlplane=not self._dry_run, copy_bundle_to=self._copy_bundle_to
+                    obj,
+                    upload_to_controlplane=not self._dry_run,
+                    copy_bundle_to=self._copy_bundle_to,
                 )
             else:
                 if self._copy_files != "none":
@@ -253,7 +261,8 @@ class _Runner:
                     pb2=run_definition_pb2.Run(
                         action=run_definition_pb2.Action(
                             id=run_definition_pb2.ActionIdentifier(
-                                name="a0", run=run_definition_pb2.RunIdentifier(name="dry-run")
+                                name="a0",
+                                run=run_definition_pb2.RunIdentifier(name="dry-run"),
                             )
                         )
                     )
@@ -286,7 +295,7 @@ class _Runner:
         if obj.parent_env is None:
             raise ValueError("Task is not attached to an environment. Please attach the task to an environment.")
 
-        image_cache = build_images.aio(cast(Environment, obj.parent_env()))
+        image_cache = await build_images.aio(cast(Environment, obj.parent_env()))
 
         code_bundle = None
         if self._name is not None:
@@ -296,7 +305,9 @@ class _Runner:
         if not code_bundle:
             if self._interactive_mode:
                 code_bundle = await build_pkl_bundle(
-                    obj, upload_to_controlplane=not self._dry_run, copy_bundle_to=self._copy_bundle_to
+                    obj,
+                    upload_to_controlplane=not self._dry_run,
+                    copy_bundle_to=self._copy_bundle_to,
                 )
             else:
                 if self._copy_files != "none":
@@ -381,7 +392,8 @@ class _Runner:
         tctx = TaskContext(
             action=action,
             checkpoints=Checkpoints(
-                prev_checkpoint_path=internal_ctx().raw_data.path, checkpoint_path=internal_ctx().raw_data.path
+                prev_checkpoint_path=internal_ctx().raw_data.path,
+                checkpoint_path=internal_ctx().raw_data.path,
             ),
             code_bundle=None,
             output_path=self._metadata_path,
@@ -403,7 +415,10 @@ class _Runner:
 
     @syncify
     async def run(
-        self, task: TaskTemplate[P, Union[R, Run]] | LazyEntity, *args: P.args, **kwargs: P.kwargs
+        self,
+        task: TaskTemplate[P, Union[R, Run]] | LazyEntity,
+        *args: P.args,
+        **kwargs: P.kwargs,
     ) -> Union[R, Run]:
         """
         Run an async `@env.task` or `TaskTemplate` instance. The existing async context will be used.
@@ -430,6 +445,10 @@ class _Runner:
 
         if isinstance(task, LazyEntity) and self._mode != "remote":
             raise ValueError("Remote task can only be run in remote mode.")
+
+        if not isinstance(task, TaskTemplate) and not isinstance(task, LazyEntity):
+            raise TypeError("On Flyte tasks can be run, not generic functions or methods.")
+
         if self._mode == "remote":
             return await self._run_remote(task, *args, **kwargs)
         task = cast(TaskTemplate, task)

flyte/_secret.py

@@ -45,6 +45,27 @@ class Secret:
             if not re.match(pattern, self.as_env_var):
                 raise ValueError(f"Invalid environment variable name: {self.as_env_var}, must match {pattern}")
 
+    def stable_hash(self) -> str:
+        """
+        Deterministic, process-independent hash (as hex string).
+        """
+        import hashlib
+
+        data = (
+            self.key,
+            self.group or "",
+            str(self.mount) if self.mount else "",
+            self.as_env_var or "",
+        )
+        joined = "|".join(data)
+        return hashlib.sha256(joined.encode("utf-8")).hexdigest()
+
+    def __hash__(self) -> int:
+        """
+        Deterministic hash function for the Secret class.
+        """
+        return int(self.stable_hash()[:16], 16)
+
 
 SecretRequest = Union[str, Secret, List[str | Secret]]
 
@@ -59,3 +80,12 @@ def secrets_from_request(secrets: SecretRequest) -> List[Secret]:
         return [secrets]
     else:
         return [Secret(key=s) if isinstance(s, str) else s for s in secrets]
+
+
+if __name__ == "__main__":
+    # Example usage
+    secret1 = Secret(key="MY_SECRET", mount=pathlib.Path("/path/to/secret"), as_env_var="MY_SECRET_ENV")
+    secret2 = Secret(
+        key="ANOTHER_SECRET",
+    )
+    print(hash(secret1), hash(secret2))

flyte/_task.py

@@ -3,7 +3,6 @@ from __future__ import annotations
 import asyncio
 import weakref
 from dataclasses import dataclass, field, replace
-from functools import cached_property
 from inspect import iscoroutinefunction
 from typing import (
     TYPE_CHECKING,
@@ -361,16 +360,13 @@ class AsyncFunctionTaskTemplate(TaskTemplate[P, R]):
     """
 
     func: FunctionTypes
+    plugin_config: Optional[Any] = None  # This is used to pass plugin specific configuration
 
     def __post_init__(self):
         super().__post_init__()
         if not iscoroutinefunction(self.func):
             self._call_as_synchronous = True
 
-    @cached_property
-    def native_interface(self) -> NativeInterface:
-        return NativeInterface.from_callable(self.func)
-
     def forward(self, *args: P.args, **kwargs: P.kwargs) -> Coroutine[Any, Any, R] | R:
         # In local execution, we want to just call the function. Note we're not awaiting anything here.
         # If the function was a coroutine function, the coroutine is returned and the await that the caller has

flyte/_task_environment.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import inspect
 import weakref
 from dataclasses import dataclass, field, replace
 from datetime import timedelta
@@ -60,13 +61,19 @@ class TaskEnvironment(Environment):
     :param reusable: Reuse policy for the environment, if set, a python process may be reused for multiple tasks.
     """
 
-    cache: Union[CacheRequest] = "auto"
+    cache: Union[CacheRequest] = "disable"
     reusable: ReusePolicy | None = None
+    plugin_config: Optional[Any] = None
     # TODO Shall we make this union of string or env? This way we can lookup the env by module/file:name
     # TODO also we could add list of files that are used by this environment
 
     _tasks: Dict[str, TaskTemplate] = field(default_factory=dict, init=False)
 
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        if self.reusable is not None and self.plugin_config is not None:
+            raise ValueError("Cannot set plugin_config when environment is reusable.")
+
     def clone_with(
         self,
         name: str,
@@ -133,6 +140,8 @@ class TaskEnvironment(Environment):
         used.
         :param report: Optional Whether to generate the html report for the task, defaults to False.
         """
+        from ._task import P, R
+
         if self.reusable is not None:
             if pod_template is not None:
                 raise ValueError("Cannot set pod_template when environment is reusable.")
@@ -141,7 +150,29 @@ class TaskEnvironment(Environment):
             friendly_name = name or func.__name__
             task_name = self.name + "." + func.__name__
 
-            tmpl: AsyncFunctionTaskTemplate = AsyncFunctionTaskTemplate(
+            if not inspect.iscoroutinefunction(func) and self.reusable is not None:
+                if self.reusable.concurrency > 1:
+                    raise ValueError(
+                        "Reusable environments with concurrency greater than 1 are only supported for async tasks. "
+                        "Please use an async function or set concurrency to 1."
+                    )
+
+            if self.plugin_config is not None:
+                from flyte.extend import TaskPluginRegistry
+
+                task_template_class: type[AsyncFunctionTaskTemplate[P, R]] | None = TaskPluginRegistry.find(
+                    config_type=type(self.plugin_config)
+                )
+                if task_template_class is None:
+                    raise ValueError(
+                        f"No task plugin found for config type {type(self.plugin_config)}. "
+                        f"Please register a plugin using flyte.extend.TaskPluginRegistry.register() api."
+                    )
+            else:
+                task_template_class = AsyncFunctionTaskTemplate[P, R]
+
+            task_template_class = cast(type[AsyncFunctionTaskTemplate[P, R]], task_template_class)
+            tmpl = task_template_class(
                 func=func,
                 name=task_name,
                 image=self.image,
@@ -158,6 +189,7 @@ class TaskEnvironment(Environment):
                 interface=NativeInterface.from_callable(func),
                 report=report,
                 friendly_name=friendly_name,
+                plugin_config=self.plugin_config,
             )
             self._tasks[task_name] = tmpl
             return tmpl

flyte/_task_plugins.py

@@ -0,0 +1,45 @@
+import typing
+from typing import Type
+
+import rich.repr
+
+if typing.TYPE_CHECKING:
+    from ._task import AsyncFunctionTaskTemplate
+
+T = typing.TypeVar("T", bound="AsyncFunctionTaskTemplate")
+
+
+class _Registry:
+    """
+    A registry for task plugins.
+    """
+
+    def __init__(self):
+        self._plugins: typing.Dict[Type, Type[T]] = {}
+
+    def register(self, config_type: Type, plugin: Type[T]):
+        """
+        Register a plugin.
+        """
+        self._plugins[config_type] = plugin
+
+    def find(self, config_type: Type) -> typing.Optional[Type[T]]:
+        """
+        Get a plugin by name.
+        """
+        return self._plugins.get(config_type)
+
+    def list_plugins(self):
+        """
+        List all registered plugins.
+        """
+        return list(self._plugins.keys())
+
+    def __rich_repr__(self) -> "rich.repr.Result":
+        yield from (("Name", i) for i in self.list_plugins())
+
+    def __repr__(self):
+        return f"TaskPluginRegistry(plugins={self.list_plugins()})"
+
+
+TaskPluginRegistry = _Registry()

flyte/_version.py

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

flyte/errors.py

@@ -157,9 +157,9 @@ class RuntimeDataValidationError(RuntimeUserError):
     This error is raised when the user tries to access a resource that does not exist or is invalid.
     """
 
-    def __init__(self, var: str, e: Exception, task_name: str = ""):
+    def __init__(self, var: str, e: Exception | str, task_name: str = ""):
         super().__init__(
-            "DataValiationError", f"In task {task_name} variable {var}, failed to serialize/deserialize because {e}"
+            "DataValiationError", f"In task {task_name} variable {var}, failed to serialize/deserialize because of {e}"
         )
 
 

flyte/extend.py

@@ -0,0 +1,12 @@
+from ._initialize import is_initialized
+from ._resources import PRIMARY_CONTAINER_DEFAULT_NAME, pod_spec_from_resources
+from ._task import AsyncFunctionTaskTemplate
+from ._task_plugins import TaskPluginRegistry
+
+__all__ = [
+    "PRIMARY_CONTAINER_DEFAULT_NAME",
+    "AsyncFunctionTaskTemplate",
+    "TaskPluginRegistry",
+    "is_initialized",
+    "pod_spec_from_resources",
+]

flyte/models.py

@@ -256,6 +256,7 @@ class NativeInterface:
     _remote_defaults: Optional[Dict[str, literals_pb2.Literal]] = field(default=None, repr=False)
 
     has_default: ClassVar[Type[_has_default]] = _has_default  # This can be used to indicate if a specific input
+
     # has a default value or not, in the case when the default value is not known. An example would be remote tasks.
 
     def has_outputs(self) -> bool:
@@ -298,7 +299,15 @@ class NativeInterface:
         sig = inspect.signature(func)
 
         # Extract parameter details (name, type, default value)
-        param_info = {name: (param.annotation, param.default) for name, param in sig.parameters.items()}
+        param_info = {}
+        for name, param in sig.parameters.items():
+            if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD):
+                raise ValueError(f"Function {func.__name__} cannot have variable positional or keyword arguments.")
+            if param.annotation is inspect.Parameter.empty:
+                logger.warning(
+                    f"Function {func.__name__} has parameter {name} without type annotation. Data will be pickled."
+                )
+            param_info[name] = (param.annotation, param.default)
 
         # Get return type
         outputs = extract_return_annotation(sig.return_annotation)

flyte/types/_type_engine.py

@@ -1315,7 +1315,8 @@ class TypeEngine(typing.Generic[T]):
             try:
                 return transformer.guess_python_type(flyte_type)
             except ValueError:
-                logger.debug(f"Skipping transformer {transformer.name} for {flyte_type}")
+                # Skipping transformer
+                continue
 
         # Because the dataclass transformer is handled explicitly in the get_transformer code, we have to handle it
         # separately here too.
@@ -1438,6 +1439,19 @@ def _type_essence(x: types_pb2.LiteralType) -> types_pb2.LiteralType:
 
 
 def _are_types_castable(upstream: types_pb2.LiteralType, downstream: types_pb2.LiteralType) -> bool:
+    if upstream.union_type is not None:
+        # for each upstream variant, there must be a compatible type downstream
+        for v in upstream.union_type.variants:
+            if not _are_types_castable(v, downstream):
+                return False
+        return True
+
+    if downstream.union_type is not None:
+        # there must be a compatible downstream type
+        for v in downstream.union_type.variants:
+            if _are_types_castable(upstream, v):
+                return True
+
     if upstream.HasField("collection_type"):
         if not downstream.HasField("collection_type"):
             return False
@@ -1483,7 +1497,7 @@ def _are_types_castable(upstream: types_pb2.LiteralType, downstream: types_pb2.L
 
         return True
 
-    if upstream.HasField("union_type"):
+    if upstream.HasField("union_type") and upstream.union_type is not None:
         # for each upstream variant, there must be a compatible type downstream
         for v in upstream.union_type.variants:
             if not _are_types_castable(v, downstream):

flyte-0.2.0b25.data/scripts/runtime.py

@@ -0,0 +1,169 @@
+"""
+Flyte runtime module, this is the entrypoint script for the Flyte runtime.
+
+Caution: Startup time for this module is very important, as it is the entrypoint for the Flyte runtime.
+Refrain from importing any modules here. If you need to import any modules, do it inside the main function.
+"""
+
+import asyncio
+import os
+import sys
+from typing import Any, List
+
+import click
+
+# Todo: work with pvditt to make these the names
+# ACTION_NAME = "_U_ACTION_NAME"
+# RUN_NAME = "_U_RUN_NAME"
+# PROJECT_NAME = "_U_PROJECT_NAME"
+# DOMAIN_NAME = "_U_DOMAIN_NAME"
+# ORG_NAME = "_U_ORG_NAME"
+
+ACTION_NAME = "ACTION_NAME"
+RUN_NAME = "RUN_NAME"
+PROJECT_NAME = "FLYTE_INTERNAL_TASK_PROJECT"
+DOMAIN_NAME = "FLYTE_INTERNAL_TASK_DOMAIN"
+ORG_NAME = "_U_ORG_NAME"
+ENDPOINT_OVERRIDE = "_U_EP_OVERRIDE"
+RUN_OUTPUT_BASE_DIR = "_U_RUN_BASE"
+ENABLE_REF_TASKS = "_REF_TASKS"  # This is a temporary flag to enable reference tasks in the runtime.
+
+# TODO: Remove this after proper auth is implemented
+_UNION_EAGER_API_KEY_ENV_VAR = "_UNION_EAGER_API_KEY"
+
+
+@click.group()
+def _pass_through():
+    pass
+
+
+@_pass_through.command("a0")
+@click.option("--inputs", "-i", required=True)
+@click.option("--outputs-path", "-o", required=True)
+@click.option("--version", "-v", required=True)
+@click.option("--run-base-dir", envvar=RUN_OUTPUT_BASE_DIR, required=True)
+@click.option("--raw-data-path", "-r", required=False)
+@click.option("--checkpoint-path", "-c", required=False)
+@click.option("--prev-checkpoint", "-p", required=False)
+@click.option("--name", envvar=ACTION_NAME, required=False)
+@click.option("--run-name", envvar=RUN_NAME, required=False)
+@click.option("--project", envvar=PROJECT_NAME, required=False)
+@click.option("--domain", envvar=DOMAIN_NAME, required=False)
+@click.option("--org", envvar=ORG_NAME, required=False)
+@click.option("--image-cache", required=False)
+@click.option("--tgz", required=False)
+@click.option("--pkl", required=False)
+@click.option("--dest", required=False)
+@click.option("--resolver", required=False)
+@click.argument(
+    "resolver-args",
+    type=click.UNPROCESSED,
+    nargs=-1,
+)
+def main(
+    run_name: str,
+    name: str,
+    project: str,
+    domain: str,
+    org: str,
+    image_cache: str,
+    version: str,
+    inputs: str,
+    run_base_dir: str,
+    outputs_path: str,
+    raw_data_path: str,
+    checkpoint_path: str,
+    prev_checkpoint: str,
+    tgz: str,
+    pkl: str,
+    dest: str,
+    resolver: str,
+    resolver_args: List[str],
+):
+    sys.path.insert(0, ".")
+
+    import flyte
+    import flyte._utils as utils
+    from flyte._initialize import init
+    from flyte._internal.controllers import create_controller
+    from flyte._internal.imagebuild.image_builder import ImageCache
+    from flyte._internal.runtime.entrypoints import load_and_run_task
+    from flyte._logging import logger
+    from flyte.models import ActionID, Checkpoints, CodeBundle, RawDataPath
+
+    logger.info(f"Initializing flyte runtime - version {flyte.__version__}")
+
+    assert org, "Org is required for now"
+    assert project, "Project is required"
+    assert domain, "Domain is required"
+    assert run_name, f"Run name is required {run_name}"
+    assert name, f"Action name is required {name}"
+
+    if run_name.startswith("{{"):
+        run_name = os.getenv("RUN_NAME", "")
+    if name.startswith("{{"):
+        name = os.getenv("ACTION_NAME", "")
+
+    # Figure out how to connect
+    # This detection of api key is a hack for now.
+    controller_kwargs: dict[str, Any] = {"insecure": False}
+    if api_key := os.getenv(_UNION_EAGER_API_KEY_ENV_VAR):
+        logger.info("Using api key from environment")
+        controller_kwargs["api_key"] = api_key
+    else:
+        ep = os.environ.get(ENDPOINT_OVERRIDE, "host.docker.internal:8090")
+        controller_kwargs["endpoint"] = ep
+        if "localhost" in ep or "docker" in ep:
+            controller_kwargs["insecure"] = True
+        logger.debug(f"Using controller endpoint: {ep} with kwargs: {controller_kwargs}")
+
+    bundle = CodeBundle(tgz=tgz, pkl=pkl, destination=dest, computed_version=version)
+    enable_ref_tasks = os.getenv(ENABLE_REF_TASKS, "false").lower() in ("true", "1", "yes")
+    # We init regular client here so that reference tasks can work
+    # Current reference tasks will not work with remote controller, because we create 2 different
+    # channels on different threads and this is not supported by grpcio or the auth system. It ends up leading
+    # File "src/python/grpcio/grpc/_cython/_cygrpc/aio/completion_queue.pyx.pxi", line 147,
+    # in grpc._cython.cygrpc.PollerCompletionQueue._handle_events
+    # BlockingIOError: [Errno 11] Resource temporarily unavailable
+    # TODO solution is to use a single channel for both controller and reference tasks, but this requires a refactor
+    if enable_ref_tasks:
+        logger.warning(
+            "Reference tasks are enabled. This will initialize client and you will see a BlockIOError. "
+            "This is harmless, but a nuisance. We are working on a fix."
+        )
+        init(org=org, project=project, domain=domain, **controller_kwargs)
+    else:
+        init()
+    # Controller is created with the same kwargs as init, so that it can be used to run tasks
+    controller = create_controller(ct="remote", **controller_kwargs)
+
+    ic = ImageCache.from_transport(image_cache) if image_cache else None
+
+    # Create a coroutine to load the task and run it
+    task_coroutine = load_and_run_task(
+        resolver=resolver,
+        resolver_args=resolver_args,
+        action=ActionID(name=name, run_name=run_name, project=project, domain=domain, org=org),
+        raw_data_path=RawDataPath(path=raw_data_path),
+        checkpoints=Checkpoints(checkpoint_path, prev_checkpoint),
+        code_bundle=bundle,
+        input_path=inputs,
+        output_path=outputs_path,
+        run_base_dir=run_base_dir,
+        version=version,
+        controller=controller,
+        image_cache=ic,
+    )
+    # Create a coroutine to watch for errors
+    controller_failure = controller.watch_for_errors()
+
+    # Run both coroutines concurrently and wait for first to finish and cancel the other
+    async def _run_and_stop():
+        await utils.run_coros(controller_failure, task_coroutine)
+        await controller.stop()
+
+    asyncio.run(_run_and_stop())
+
+
+if __name__ == "__main__":
+    _pass_through()

{flyte-0.2.0b23.dist-info → flyte-0.2.0b25.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flyte
-Version: 0.2.0b23
+Version: 0.2.0b25
 Summary: Add your description here
 Author-email: Ketan Umare <kumare3@users.noreply.github.com>
 Requires-Python: >=3.10