flyte 0.2.0b5__py3-none-any.whl → 0.2.0b8__py3-none-any.whl

flyte/__init__.py

@@ -38,6 +38,7 @@ __all__ = [
     "deploy",
     "group",
     "init",
+    "init_auto_from_config",
     "run",
     "trace",
     "with_runcontext",
@@ -49,7 +50,7 @@ from ._deploy import deploy
 from ._environment import Environment
 from ._group import group
 from ._image import Image
-from ._initialize import init
+from ._initialize import init, init_auto_from_config
 from ._resources import GPU, TPU, Device, Resources
 from ._retry import RetryStrategy
 from ._reusable_environment import ReusePolicy

flyte/_code_bundle/_utils.py

@@ -14,7 +14,6 @@ import tempfile
 import typing
 from datetime import datetime, timezone
 from functools import lru_cache
-from pathlib import Path
 from types import ModuleType
 from typing import List, Literal, Optional, Tuple, Union
 
@@ -322,18 +321,3 @@ def hash_file(file_path: typing.Union[os.PathLike, str]) -> Tuple[bytes, str, in
             size += len(chunk)
 
     return h.digest(), h.hexdigest(), size
-
-
-def _find_project_root(source_path) -> str:
-    """
-    Find the root of the project.
-    The root of the project is considered to be the first ancestor from source_path that does
-    not contain a __init__.py file.
-
-    N.B.: This assumption only holds for regular packages (as opposed to namespace packages)
-    """
-    # Start from the directory right above source_path
-    path = Path(source_path).parent.resolve()
-    while os.path.exists(os.path.join(path, "__init__.py")):
-        path = path.parent
-    return str(path)

flyte/_code_bundle/bundle.py

@@ -146,7 +146,7 @@ async def build_code_bundle(
         logger.info(f"Code bundle created at {bundle_path}, size: {tar_size} MB, archive size: {archive_size} MB")
         if not dryrun:
             hash_digest, remote_path = await upload_file(bundle_path)
-            logger.info(f"Code bundle uploaded to {remote_path}")
+            logger.debug(f"Code bundle uploaded to {remote_path}")
         else:
             remote_path = "na"
             if copy_bundle_to:

flyte/_environment.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
+import re
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union
 
 import rich.repr
 
@@ -14,6 +15,10 @@ if TYPE_CHECKING:
     from kubernetes.client import V1PodTemplate
 
 
+def is_snake_or_kebab_with_numbers(s: str) -> bool:
+    return re.fullmatch(r"^[a-z0-9]+([_-][a-z0-9]+)*$", s) is not None
+
+
 @rich.repr.auto
 @dataclass(init=True, repr=True)
 class Environment:
@@ -36,8 +41,44 @@ class Environment:
     resources: Optional[Resources] = None
     image: Union[str, Image, Literal["auto"]] = "auto"
 
+    def __post_init__(self):
+        if not is_snake_or_kebab_with_numbers(self.name):
+            raise ValueError(f"Environment name '{self.name}' must be in snake_case or kebab-case format.")
+
     def add_dependency(self, *env: Environment):
         """
         Add a dependency to the environment.
         """
         self.env_dep_hints.extend(env)
+
+    def clone_with(
+        self,
+        name: str,
+        image: Optional[Union[str, Image, Literal["auto"]]] = None,
+        resources: Optional[Resources] = None,
+        env: Optional[Dict[str, str]] = None,
+        secrets: Optional[SecretRequest] = None,
+        env_dep_hints: Optional[List[Environment]] = None,
+        **kwargs: Any,
+    ) -> Environment:
+        raise NotImplementedError
+
+    def _get_kwargs(self) -> Dict[str, Any]:
+        """
+        Get the keyword arguments for the environment.
+        """
+        kwargs: Dict[str, Any] = {
+            "env_dep_hints": self.env_dep_hints,
+            "image": self.image,
+        }
+        if self.resources is not None:
+            kwargs["resources"] = self.resources
+        if self.secrets is not None:
+            kwargs["secrets"] = self.secrets
+        if self.env is not None:
+            kwargs["env"] = self.env
+        if self.pod_template is not None:
+            kwargs["pod_template"] = self.pod_template
+        if self.description is not None:
+            kwargs["description"] = self.description
+        return kwargs

flyte/_image.py

@@ -444,8 +444,7 @@ class Image:
         ```
 
         For more information on the uv script format, see the documentation:
-        <href="https://docs.astral.sh/uv/guides/scripts/#declaring-script-dependencies">
-            UV: Declaring script dependencies</href>
+        [UV: Declaring script dependencies](https://docs.astral.sh/uv/guides/scripts/#declaring-script-dependencies)
 
         :param name: name of the image
         :param registry: registry to use for the image

flyte/_initialize.py

@@ -10,7 +10,7 @@ from typing import TYPE_CHECKING, Callable, List, Literal, Optional, TypeVar
 from flyte.errors import InitializationError
 from flyte.syncify import syncify
 
-from ._logging import initialize_logger
+from ._logging import initialize_logger, logger
 from ._tools import ipython_check
 
 if TYPE_CHECKING:
@@ -130,7 +130,6 @@ async def init(
     rpc_retries: int = 3,
     http_proxy_url: str | None = None,
     storage: Storage | None = None,
-    config: Config | None = None,
 ) -> None:
     """
     Initialize the Flyte system with the given configuration. This method should be called before any other Flyte
@@ -179,42 +178,72 @@ async def init(
     global _init_config  # noqa: PLW0603
 
     with _init_lock:
-        if config is None:
-            import flyte.config as _f_cfg
-
-            config = _f_cfg.Config()
-        platform_cfg = config.platform
-        task_cfg = config.task
         client = None
-        if endpoint or platform_cfg.endpoint or api_key:
+        if endpoint or api_key:
             client = await _initialize_client(
                 api_key=api_key,
-                auth_type=auth_type or platform_cfg.auth_mode,
-                endpoint=endpoint or platform_cfg.endpoint,
+                auth_type=auth_type,
+                endpoint=endpoint,
                 headless=headless,
-                insecure=insecure or platform_cfg.insecure,
-                insecure_skip_verify=insecure_skip_verify or platform_cfg.insecure_skip_verify,
-                ca_cert_file_path=ca_cert_file_path or platform_cfg.ca_cert_file_path,
-                command=command or platform_cfg.command,
-                proxy_command=proxy_command or platform_cfg.proxy_command,
-                client_id=client_id or platform_cfg.client_id,
-                client_credentials_secret=client_credentials_secret or platform_cfg.client_credentials_secret,
+                insecure=insecure,
+                insecure_skip_verify=insecure_skip_verify,
+                ca_cert_file_path=ca_cert_file_path,
+                command=command,
+                proxy_command=proxy_command,
+                client_id=client_id,
+                client_credentials_secret=client_credentials_secret,
                 client_config=auth_client_config,
-                rpc_retries=rpc_retries or platform_cfg.rpc_retries,
-                http_proxy_url=http_proxy_url or platform_cfg.http_proxy_url,
+                rpc_retries=rpc_retries,
+                http_proxy_url=http_proxy_url,
             )
 
         root_dir = root_dir or get_cwd_editable_install() or Path.cwd()
         _init_config = _InitConfig(
             root_dir=root_dir,
-            project=project or task_cfg.project,
-            domain=domain or task_cfg.domain,
+            project=project,
+            domain=domain,
             client=client,
             storage=storage,
-            org=org or task_cfg.org,
+            org=org,
         )
 
 
+@syncify
+async def init_auto_from_config(path_or_config: str | Config | None = None) -> None:
+    """
+    Initialize the Flyte system using a configuration file or Config object. This method should be called before any
+    other Flyte remote API methods are called. Thread-safe implementation.
+
+    :param path_or_config: Path to the configuration file or Config object
+    :return: None
+    """
+    import flyte.config as config
+
+    cfg: config.Config
+    if path_or_config is None or isinstance(path_or_config, str):
+        # If a string is passed, treat it as a path to the config file
+        cfg = config.auto(path_or_config)
+    else:
+        # If a Config object is passed, use it directly
+        cfg = path_or_config
+
+    logger.debug(f"Flyte config initialized as {cfg}")
+    await init.aio(
+        org=cfg.task.org,
+        project=cfg.task.project,
+        domain=cfg.task.domain,
+        endpoint=cfg.platform.endpoint,
+        insecure=cfg.platform.insecure,
+        insecure_skip_verify=cfg.platform.insecure_skip_verify,
+        ca_cert_file_path=cfg.platform.ca_cert_file_path,
+        auth_type=cfg.platform.auth_mode,
+        command=cfg.platform.command,
+        proxy_command=cfg.platform.proxy_command,
+        client_id=cfg.platform.client_id,
+        client_credentials_secret=cfg.platform.client_credentials_secret,
+    )
+
+
 def _get_init_config() -> Optional[_InitConfig]:
     """
     Get the current initialization configuration. Thread-safe implementation.

flyte/_internal/controllers/__init__.py

@@ -28,6 +28,8 @@ class Controller(Protocol):
         """
         ...
 
+    def submit_sync(self, _task: TaskTemplate, *args, **kwargs) -> Any: ...
+
     async def submit_task_ref(self, _task: task_definition_pb2.TaskDetails, *args, **kwargs) -> Any:
         """
         Submit a task reference to the controller asynchronously and wait for the result. This is async and will block

flyte/_internal/controllers/_local_controller.py

@@ -8,6 +8,7 @@ from flyte._internal.runtime.entrypoints import direct_dispatch
 from flyte._logging import log, logger
 from flyte._protos.workflow import task_definition_pb2
 from flyte._task import TaskTemplate
+from flyte._utils.asyn import loop_manager
 from flyte.models import ActionID, NativeInterface, RawDataPath
 
 R = TypeVar("R")
@@ -58,6 +59,8 @@ class LocalController:
             return result
         return out
 
+    submit_sync = loop_manager.synced(submit)
+
     async def finalize_parent_action(self, action: ActionID):
         pass
 

flyte/_internal/controllers/remote/_controller.py

@@ -21,6 +21,7 @@ from flyte._internal.runtime.task_serde import translate_task_to_wire
 from flyte._logging import logger
 from flyte._protos.workflow import run_definition_pb2, task_definition_pb2
 from flyte._task import TaskTemplate
+from flyte._utils.asyn import loop_manager
 from flyte.models import ActionID, NativeInterface, SerializationContext
 
 R = TypeVar("R")
@@ -235,6 +236,8 @@ class RemoteController(Controller):
         async with self._parent_action_semaphore[unique_action_name(current_action_id)]:
             return await self._submit(task_call_seq, _task, *args, **kwargs)
 
+    submit_sync = loop_manager.synced(submit)
+
     async def finalize_parent_action(self, action_id: ActionID):
         """
         This method is invoked when the parent action is finished. It will finalize the run and upload the outputs

flyte/_internal/controllers/remote/_core.py

@@ -32,7 +32,7 @@ class Controller:
         max_system_retries: int = 5,
         resource_log_interval_sec: float = 10.0,
         min_backoff_on_err_sec: float = 0.1,
-        thread_wait_timeout_sec: float = 0.5,
+        thread_wait_timeout_sec: float = 5.0,
         enqueue_timeout_sec: float = 5.0,
     ):
         """

flyte/_internal/controllers/remote/_informer.py

@@ -38,11 +38,11 @@ class ActionCache:
         """
         Add an action to the cache if it doesn't exist. This is invoked by the watch.
         """
-        logger.info(f"Observing phase {run_definition_pb2.Phase.Name(state.phase)} for {state.action_id.name}")
+        logger.debug(f"Observing phase {run_definition_pb2.Phase.Name(state.phase)} for {state.action_id.name}")
         if state.output_uri:
-            logger.info(f"Output URI: {state.output_uri}")
+            logger.debug(f"Output URI: {state.output_uri}")
         else:
-            logger.info(f"{state.action_id.name} has no output URI")
+            logger.warning(f"{state.action_id.name} has no output URI")
         if state.phase == run_definition_pb2.Phase.PHASE_FAILED:
             logger.error(
                 f"Action {state.action_id.name} failed with error (msg):"

flyte/_task.py

@@ -3,10 +3,10 @@ from __future__ import annotations
 import weakref
 from dataclasses import dataclass, field, replace
 from functools import cached_property
+from inspect import iscoroutinefunction
 from typing import (
     TYPE_CHECKING,
     Any,
-    Awaitable,
     Callable,
     Coroutine,
     Dict,
@@ -15,6 +15,7 @@ from typing import (
     Literal,
     Optional,
     ParamSpec,
+    TypeAlias,
     TypeVar,
     Union,
 )
@@ -42,6 +43,10 @@ if TYPE_CHECKING:
 P = ParamSpec("P")  # capture the function's parameters
 R = TypeVar("R")  # return type
 
+AsyncFunctionType: TypeAlias = Callable[P, Coroutine[Any, Any, R]]
+SyncFunctionType: TypeAlias = Callable[P, R]
+FunctionTypes: TypeAlias = Union[AsyncFunctionType, SyncFunctionType]
+
 
 @dataclass(kw_only=True)
 class TaskTemplate(Generic[P, R]):
@@ -98,6 +103,9 @@ class TaskTemplate(Generic[P, R]):
     local: bool = field(default=False, init=False)
     ref: bool = field(default=False, init=False, repr=False, compare=False)
 
+    # Only used in python 3.10 and 3.11, where we cannot use markcoroutinefunction
+    _call_as_synchronous: bool = False
+
     def __post_init__(self):
         # If pod_template is set to a pod, verify
         if self.pod_template is not None and not isinstance(self.pod_template, str):
@@ -208,13 +216,15 @@ class TaskTemplate(Generic[P, R]):
     def native_interface(self) -> NativeInterface:
         return self.interface
 
-    async def __call__(self, *args: P.args, **kwargs: P.kwargs) -> Coroutine[Any, Any, R]:
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> Coroutine[Any, Any, R] | R:
         """
         This is the entrypoint for an async function task at runtime. It will be called during an execution.
         Please do not override this method, if you simply want to modify the execution behavior, override the
         execute method.
 
-        # TODO lets provide one hook to implement, _pre, _execute and _post. We do not want actual execute to
+        This needs to be overridable to maybe be async.
+        The returned thing from here needs to be an awaitable if the underlying task is async, and a regular object
+        if the task is not.
         """
         try:
             ctx = internal_ctx()
@@ -226,8 +236,14 @@ class TaskTemplate(Generic[P, R]):
 
                 controller = get_controller()
                 if controller:
-                    return await controller.submit(self, *args, **kwargs)
-            return await self.execute(*args, **kwargs)
+                    if self._call_as_synchronous:
+                        return controller.submit_sync(self, *args, **kwargs)
+                    else:
+                        return controller.submit(self, *args, **kwargs)
+                else:
+                    raise RuntimeSystemError("BadContext", "Controller is not initialized.")
+
+            return self.forward(*args, **kwargs)
         except RuntimeSystemError:
             raise
         except RuntimeUserError:
@@ -235,6 +251,17 @@ class TaskTemplate(Generic[P, R]):
         except Exception as e:
             raise RuntimeUserError(type(e).__name__, str(e)) from e
 
+    def forward(self, *args: P.args, **kwargs: P.kwargs) -> Coroutine[Any, Any, R] | R:
+        """
+        Think of this as a local execute method for your task. This function will be invoked by the __call__ method
+        when not in a Flyte task execution context.  See the implementation below for an example.
+
+        :param args:
+        :param kwargs:
+        :return:
+        """
+        raise NotImplementedError
+
     def override(
         self,
         *,
@@ -290,26 +317,38 @@ class AsyncFunctionTaskTemplate(TaskTemplate[P, R]):
     is decorated with the task decorator.
     """
 
-    func: Callable[P, Awaitable[R]]
+    func: FunctionTypes
+
+    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
+        # in front of the task invocation will handle the awaiting.
+        return self.func(*args, **kwargs)
+
     async def execute(self, *args: P.args, **kwargs: P.kwargs) -> R:
         """
         This is the execute method that will be called when the task is invoked. It will call the actual function.
         # TODO We may need to keep this as the bare func execute, and need a pre and post execute some other func.
         """
+
         ctx = internal_ctx()
+        assert ctx.data.task_context is not None, "Function should have already returned if not in a task context"
         ctx_data = await self.pre(*args, **kwargs)
-        if ctx.data.task_context is not None:
-            tctx = ctx.data.task_context.replace(data=ctx_data)
-            with ctx.replace_task_context(tctx):
+        tctx = ctx.data.task_context.replace(data=ctx_data)
+        with ctx.replace_task_context(tctx):
+            if iscoroutinefunction(self.func):
                 v = await self.func(*args, **kwargs)
-                await self.post(v)
-        else:
-            v = await self.func(*args, **kwargs)
+            else:
+                v = self.func(*args, **kwargs)
             await self.post(v)
         return v
 

flyte/_task_environment.py

@@ -1,11 +1,19 @@
 from __future__ import annotations
 
-import asyncio
 import weakref
 from dataclasses import dataclass, field, replace
 from datetime import timedelta
-from functools import wraps
-from typing import TYPE_CHECKING, Awaitable, Callable, Dict, List, Literal, Optional, ParamSpec, TypeVar, Union, cast
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Union,
+    cast,
+)
 
 import rich.repr
 
@@ -23,8 +31,7 @@ from .models import NativeInterface
 if TYPE_CHECKING:
     from kubernetes.client import V1PodTemplate
 
-P = ParamSpec("P")  # capture the function's parameters
-R = TypeVar("R")  # return type
+    from ._task import FunctionTypes, P, R
 
 
 @rich.repr.auto
@@ -54,7 +61,7 @@ class TaskEnvironment(Environment):
     """
 
     cache: Union[CacheRequest] = "auto"
-    reusable: Union[ReusePolicy, Literal["auto"], None] = None
+    reusable: ReusePolicy | None = 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
 
@@ -65,32 +72,42 @@ class TaskEnvironment(Environment):
         name: str,
         image: Optional[Union[str, Image, Literal["auto"]]] = None,
         resources: Optional[Resources] = None,
-        cache: Union[CacheRequest, None] = None,
         env: Optional[Dict[str, str]] = None,
-        reusable: Union[ReusePolicy, None] = None,
         secrets: Optional[SecretRequest] = None,
         env_dep_hints: Optional[List[Environment]] = None,
+        **kwargs: Any,
     ) -> TaskEnvironment:
         """
-        Clone the environment with new settings.
+        Clone the TaskEnvironment with new parameters.
+        besides the base environment parameters, you can override, kwargs like `cache`, `reusable`, etc.
+
         """
-        if image is None:
-            image = self.image
-        else:
-            image = "auto"
-        return replace(
-            self,
-            cache=cache if cache else "auto",
-            reusable=reusable,
-            name=name,
-            image=image,
-            resources=resources,
-            env=env,
-            secrets=secrets,
-            env_dep_hints=env_dep_hints if env_dep_hints else [],
-        )
-
-    def _task(
+        cache = kwargs.pop("cache", None)
+        reusable = kwargs.pop("reusable", None)
+
+        # validate unknown kwargs if needed
+        if kwargs:
+            raise TypeError(f"Unexpected keyword arguments: {list(kwargs.keys())}")
+
+        kwargs = self._get_kwargs()
+        kwargs["name"] = name
+        if image is not None:
+            kwargs["image"] = image
+        if resources is not None:
+            kwargs["resources"] = resources
+        if cache is not None:
+            kwargs["cache"] = cache
+        if env is not None:
+            kwargs["env"] = env
+        if reusable is not None:
+            kwargs["reusable"] = reusable
+        if secrets is not None:
+            kwargs["secrets"] = secrets
+        if env_dep_hints is not None:
+            kwargs["env_dep_hints"] = env_dep_hints
+        return replace(self, **kwargs)
+
+    def task(
         self,
         _func=None,
         *,
@@ -104,6 +121,7 @@ class TaskEnvironment(Environment):
         report: bool = False,
     ) -> Union[AsyncFunctionTaskTemplate, Callable[P, R]]:
         """
+        :param _func: Optional The function to decorate. If not provided, the decorator will return a callable that
         :param name: Optional The name of the task (defaults to the function name)
         :param cache: Optional The cache policy for the task, defaults to auto, which will cache the results of the
         task.
@@ -119,25 +137,12 @@ class TaskEnvironment(Environment):
             if pod_template is not None:
                 raise ValueError("Cannot set pod_template when environment is reusable.")
 
-        def decorator(func: Callable[P, Awaitable[R]]) -> AsyncFunctionTaskTemplate[P, R]:
+        def decorator(func: FunctionTypes) -> AsyncFunctionTaskTemplate[P, R]:
             task_name = name or func.__name__
             task_name = self.name + "." + task_name
-            if len(task_name) > 30:
-                # delete this if we can remove the restriction that task names must be <= 30 characters
-                task_name = task_name[-30:]
-
-            @wraps(func)
-            async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
-                return await func(*args, **kwargs)
-
-            if not asyncio.iscoroutinefunction(func):
-                raise TypeError(
-                    f"Function {func.__name__} is not a coroutine function. Use @env.task decorator for async tasks."
-                    f"You can simply mark your function as async def {func.__name__} to make it a coroutine function, "
-                    f"it is ok to write sync code in async functions, but not the other way around."
-                )
-            tmpl = AsyncFunctionTaskTemplate(
-                func=wrapper,
+
+            tmpl: AsyncFunctionTaskTemplate = AsyncFunctionTaskTemplate(
+                func=func,
                 name=task_name,
                 image=self.image,
                 resources=self.resources,
@@ -160,29 +165,6 @@ class TaskEnvironment(Environment):
             return cast(AsyncFunctionTaskTemplate, decorator)
         return cast(AsyncFunctionTaskTemplate, decorator(_func))
 
-    @property
-    def task(self) -> Callable:
-        """
-        Decorator to create a new task with the environment settings.
-        The task will be executed in its own container with the specified image, resources, and environment variables,
-        unless reusePolicy is set, in which case the same container will be reused for all tasks with the same
-        environment settings.
-
-        :param name: Optional The name of the task (defaults to the function name)
-        :param cache: Optional The cache policy for the task, defaults to auto, which will cache the results of the
-         task.
-        :param retries: Optional The number of retries for the task, defaults to 0, which means no retries.
-        :param docs: Optional The documentation for the task, if not provided the function docstring will be used.
-        :param secrets: Optional The secrets that will be injected into the task at runtime.
-        :param timeout: Optional The timeout for the task.
-        :param pod_template: Optional The pod template for the task, if not provided the default pod template will be
-         used.
-        :param report: Optional Whether to generate the html report for the task, defaults to False.
-
-        :return: New Task instance or Task decorator
-        """
-        return self._task
-
     @property
     def tasks(self) -> Dict[str, TaskTemplate]:
         """

flyte/_utils/coro_management.py

@@ -21,5 +21,3 @@ async def run_coros(*coros: typing.Coroutine, return_when: str = asyncio.FIRST_C
         err = t.exception()
         if err:
             raise err
-        else:
-            print(f"Task result: {t.result()}")

flyte/_version.py

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