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

flyte/remote/_task.py

@@ -1,24 +1,23 @@
 from __future__ import annotations
 
+import asyncio
 import functools
 from dataclasses import dataclass
-from threading import Lock
 from typing import Any, AsyncIterator, Callable, Coroutine, Dict, Iterator, Literal, Optional, Tuple, Union, cast
 
 import rich.repr
-from flyteidl.core import literals_pb2
-from google.protobuf import timestamp
+from flyteidl2.common import identifier_pb2, list_pb2
+from flyteidl2.core import literals_pb2
+from flyteidl2.task import task_definition_pb2, task_service_pb2
 
 import flyte
 import flyte.errors
 from flyte._cache.cache import CacheBehavior
 from flyte._context import internal_ctx
-from flyte._initialize import ensure_client, get_client, get_common_config
+from flyte._initialize import ensure_client, get_client, get_init_config
 from flyte._internal.runtime.resources_serde import get_proto_resources
 from flyte._internal.runtime.task_serde import get_proto_retry_strategy, get_proto_timeout, get_security_context
 from flyte._logging import logger
-from flyte._protos.common import identifier_pb2, list_pb2
-from flyte._protos.workflow import task_definition_pb2, task_service_pb2
 from flyte.models import NativeInterface
 from flyte.syncify import syncify
 
@@ -35,7 +34,7 @@ def _repr_task_metadata(metadata: task_definition_pb2.TaskMetadata) -> rich.repr
         else:
             yield "deployed_by", f"App: {metadata.deployed_by.application.spec.name}"
     yield "short_name", metadata.short_name
-    yield "deployed_at", timestamp.to_datetime(metadata.deployed_at)
+    yield "deployed_at", metadata.deployed_at.ToDatetime()
     yield "environment_name", metadata.environment_name
 
 
@@ -49,7 +48,7 @@ class LazyEntity:
         self._task: Optional[TaskDetails] = None
         self._getter = getter
         self._name = name
-        self._mutex = Lock()
+        self._mutex = asyncio.Lock()
 
     @property
     def name(self) -> str:
@@ -60,11 +59,11 @@ class LazyEntity:
         """
         Forwards all other attributes to task, causing the task to be fetched!
         """
-        with self._mutex:
+        async with self._mutex:
             if self._task is None:
                 self._task = await self._getter()
-            if self._task is None:
-                raise RuntimeError(f"Error downloading the task {self._name}, (check original exception...)")
+                if self._task is None:
+                    raise RuntimeError(f"Error downloading the task {self._name}, (check original exception...)")
             return self._task
 
     @syncify
@@ -73,8 +72,10 @@ class LazyEntity:
         **kwargs: Any,
     ) -> LazyEntity:
         task_details = cast(TaskDetails, await self.fetch.aio())
-        task_details.override(**kwargs)
-        return self
+        new_task_details = task_details.override(**kwargs)
+        new_entity = LazyEntity(self._name, self._getter)
+        new_entity._task = new_task_details
+        return new_entity
 
     async def __call__(self, *args, **kwargs):
         """
@@ -93,10 +94,11 @@ class LazyEntity:
 AutoVersioning = Literal["latest", "current"]
 
 
-@dataclass
+@dataclass(frozen=True)
 class TaskDetails(ToJSONMixin):
     pb2: task_definition_pb2.TaskDetails
     max_inline_io_bytes: int = 10 * 1024 * 1024  # 10 MB
+    overriden_queue: Optional[str] = None
 
     @classmethod
     def get(
@@ -148,7 +150,7 @@ class TaskDetails(ToJSONMixin):
                     if ctx is None:
                         raise ValueError("auto_version=current can only be used within a task context.")
                     _version = ctx.version
-            cfg = get_common_config()
+            cfg = get_init_config()
             task_id = task_definition_pb2.TaskIdentifier(
                 org=cfg.org,
                 project=project or cfg.project,
@@ -261,12 +263,6 @@ class TaskDetails(ToJSONMixin):
                 f"Reference task {self.name} does not support positional arguments"
                 f"currently. Please use keyword arguments."
             )
-        if len(self.required_args) > 0:
-            if len(args) + len(kwargs) < len(self.required_args):
-                raise ValueError(
-                    f"Task {self.name} requires at least {self.required_args} arguments, "
-                    f"but only received args:{args}  kwargs{kwargs}."
-                )
 
         ctx = internal_ctx()
         if ctx.is_task_context():
@@ -276,19 +272,37 @@ class TaskDetails(ToJSONMixin):
             from flyte._internal.controllers import get_controller
 
             controller = get_controller()
+            if len(self.required_args) > 0:
+                if len(args) + len(kwargs) < len(self.required_args):
+                    raise ValueError(
+                        f"Task {self.name} requires at least {self.required_args} arguments, "
+                        f"but only received args:{args}  kwargs{kwargs}."
+                    )
             if controller:
-                return await controller.submit_task_ref(self.pb2, self.max_inline_io_bytes, *args, **kwargs)
-        raise flyte.errors
+                return await controller.submit_task_ref(self, *args, **kwargs)
+        raise flyte.errors.ReferenceTaskError(
+            f"Reference tasks [{self.name}] cannot be executed locally, only remotely."
+        )
+
+    @property
+    def queue(self) -> Optional[str]:
+        """
+        The queue to use for the task.
+        """
+        return self.overriden_queue
 
     def override(
         self,
         *,
-        friendly_name: Optional[str] = None,
+        short_name: Optional[str] = None,
         resources: Optional[flyte.Resources] = None,
         retries: Union[int, flyte.RetryStrategy] = 0,
         timeout: Optional[flyte.TimeoutType] = None,
         env_vars: Optional[Dict[str, str]] = None,
         secrets: Optional[flyte.SecretRequest] = None,
+        max_inline_io_bytes: Optional[int] = None,
+        cache: Optional[flyte.Cache] = None,
+        queue: Optional[str] = None,
         **kwargs: Any,
     ) -> TaskDetails:
         if len(kwargs) > 0:
@@ -296,29 +310,57 @@ class TaskDetails(ToJSONMixin):
                 f"ReferenceTasks [{self.name}] do not support overriding with kwargs: {kwargs}, "
                 f"Check the parameters for override method."
             )
-        template = self.pb2.spec.task_template
-        if friendly_name:
-            self.pb2.metadata.short_name = friendly_name
+        pb2 = task_definition_pb2.TaskDetails()
+        pb2.CopyFrom(self.pb2)
+
+        if short_name:
+            pb2.metadata.short_name = short_name
+
+        template = pb2.spec.task_template
         if secrets:
             template.security_context.CopyFrom(get_security_context(secrets))
+
         if template.HasField("container"):
             if env_vars:
                 template.container.env.clear()
                 template.container.env.extend([literals_pb2.KeyValuePair(key=k, value=v) for k, v in env_vars.items()])
             if resources:
                 template.container.resources.CopyFrom(get_proto_resources(resources))
+
+        md = template.metadata
         if retries:
-            template.metadata.retries.CopyFrom(get_proto_retry_strategy(retries))
-        if timeout:
-            template.metadata.timeout.CopyFrom(get_proto_timeout(timeout))
+            md.retries.CopyFrom(get_proto_retry_strategy(retries))
 
-        return self
+        if timeout:
+            md.timeout.CopyFrom(get_proto_timeout(timeout))
+
+        if cache:
+            if cache.behavior == "disable":
+                md.discoverable = False
+                md.discovery_version = ""
+            elif cache.behavior == "override":
+                md.discoverable = True
+                if not cache.version_override:
+                    raise ValueError("cache.version_override must be set when cache.behavior is 'override'")
+                md.discovery_version = cache.version_override
+            else:
+                if cache.behavior == "auto":
+                    raise ValueError("cache.behavior must be 'disable' or 'override' for reference tasks")
+                raise ValueError(f"Invalid cache behavior: {cache.behavior}.")
+            md.cache_serializable = cache.serialize
+            md.cache_ignore_input_vars[:] = list(cache.ignored_inputs or ())
+
+        return TaskDetails(
+            pb2,
+            max_inline_io_bytes=max_inline_io_bytes or self.max_inline_io_bytes,
+            overriden_queue=queue,
+        )
 
     def __rich_repr__(self) -> rich.repr.Result:
         """
         Rich representation of the task.
         """
-        yield "friendly_name", self.pb2.spec.short_name
+        yield "short_name", self.pb2.spec.short_name
         yield "environment", self.pb2.spec.environment
         yield "default_inputs_keys", self.default_input_args
         yield "required_args", self.required_args
@@ -408,7 +450,7 @@ class Task(ToJSONMixin):
         sort_pb2 = list_pb2.Sort(
             key=sort_by[0], direction=list_pb2.Sort.ASCENDING if sort_by[1] == "asc" else list_pb2.Sort.DESCENDING
         )
-        cfg = get_common_config()
+        cfg = get_init_config()
         filters = []
         if by_task_name:
             filters.append(

flyte/remote/_trigger.py

@@ -0,0 +1,306 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from functools import cached_property
+from typing import AsyncIterator
+
+import grpc.aio
+from flyteidl2.common import identifier_pb2, list_pb2
+from flyteidl2.task import common_pb2, task_definition_pb2
+from flyteidl2.trigger import trigger_definition_pb2, trigger_service_pb2
+
+import flyte
+from flyte._initialize import ensure_client, get_client, get_init_config
+from flyte._internal.runtime import trigger_serde
+from flyte.syncify import syncify
+
+from ._common import ToJSONMixin
+from ._task import Task, TaskDetails
+
+
+@dataclass
+class TriggerDetails(ToJSONMixin):
+    pb2: trigger_definition_pb2.TriggerDetails
+
+    @syncify
+    @classmethod
+    async def get(cls, *, name: str, task_name: str) -> TriggerDetails:
+        """
+        Retrieve detailed information about a specific trigger by its name.
+        """
+        ensure_client()
+        cfg = get_init_config()
+        resp = await get_client().trigger_service.GetTriggerDetails(
+            request=trigger_service_pb2.GetTriggerDetailsRequest(
+                name=identifier_pb2.TriggerName(
+                    task_name=task_name,
+                    name=name,
+                    org=cfg.org,
+                    project=cfg.project,
+                    domain=cfg.domain,
+                ),
+            )
+        )
+        return cls(pb2=resp.trigger)
+
+    @property
+    def name(self) -> str:
+        return self.id.name.name
+
+    @property
+    def id(self) -> identifier_pb2.TriggerIdentifier:
+        return self.pb2.id
+
+    @property
+    def task_name(self) -> str:
+        return self.pb2.id.name.task_name
+
+    @property
+    def automation_spec(self) -> common_pb2.TriggerAutomationSpec:
+        return self.pb2.automation_spec
+
+    @property
+    def metadata(self) -> trigger_definition_pb2.TriggerMetadata:
+        return self.pb2.metadata
+
+    @property
+    def status(self) -> trigger_definition_pb2.TriggerStatus:
+        return self.pb2.status
+
+    @property
+    def is_active(self) -> bool:
+        return self.pb2.spec.active
+
+    @cached_property
+    def trigger(self) -> trigger_definition_pb2.Trigger:
+        return trigger_definition_pb2.Trigger(
+            id=self.pb2.id,
+            automation_spec=self.automation_spec,
+            metadata=self.metadata,
+            status=self.status,
+            active=self.is_active,
+        )
+
+
+@dataclass
+class Trigger(ToJSONMixin):
+    pb2: trigger_definition_pb2.Trigger
+    details: TriggerDetails | None = None
+
+    @syncify
+    @classmethod
+    async def create(
+        cls,
+        trigger: flyte.Trigger,
+        task_name: str,
+        task_version: str | None = None,
+    ) -> Trigger:
+        """
+        Create a new trigger in the Flyte platform.
+
+        :param trigger: The flyte.Trigger object containing the trigger definition.
+        :param task_name: Optional name of the task to associate with the trigger.
+        """
+        ensure_client()
+        cfg = get_init_config()
+
+        # Fetch the task to ensure it exists and to get its input definitions
+        try:
+            lazy = (
+                Task.get(name=task_name, version=task_version)
+                if task_version
+                else Task.get(name=task_name, auto_version="latest")
+            )
+            task: TaskDetails = await lazy.fetch.aio()
+
+            task_trigger = await trigger_serde.to_task_trigger(
+                t=trigger,
+                task_name=task_name,
+                task_inputs=task.pb2.spec.task_template.interface.inputs,
+                task_default_inputs=list(task.pb2.spec.default_inputs),
+            )
+
+            resp = await get_client().trigger_service.DeployTrigger(
+                request=trigger_service_pb2.DeployTriggerRequest(
+                    name=identifier_pb2.TriggerName(
+                        name=trigger.name,
+                        task_name=task_name,
+                        org=cfg.org,
+                        project=cfg.project,
+                        domain=cfg.domain,
+                    ),
+                    spec=trigger_definition_pb2.TriggerSpec(
+                        active=task_trigger.spec.active,
+                        inputs=task_trigger.spec.inputs,
+                        run_spec=task_trigger.spec.run_spec,
+                        task_version=task.version,
+                    ),
+                    automation_spec=task_trigger.automation_spec,
+                )
+            )
+
+            details = TriggerDetails(pb2=resp.trigger)
+
+            return cls(pb2=details.trigger, details=details)
+        except grpc.aio.AioRpcError as e:
+            if e.code() == grpc.StatusCode.NOT_FOUND:
+                raise ValueError(f"Task {task_name}:{task_version or 'latest'} not found") from e
+            raise
+
+    @syncify
+    @classmethod
+    async def get(cls, *, name: str, task_name: str) -> TriggerDetails:
+        """
+        Retrieve a trigger by its name and associated task name.
+        """
+        return await TriggerDetails.get.aio(name=name, task_name=task_name)
+
+    @syncify
+    @classmethod
+    async def listall(
+        cls, task_name: str | None = None, task_version: str | None = None, limit: int = 100
+    ) -> AsyncIterator[Trigger]:
+        """
+        List all triggers associated with a specific task or all tasks if no task name is provided.
+        """
+        ensure_client()
+        cfg = get_init_config()
+        token = None
+        task_name_id = None
+        project_id = None
+        task_id = None
+        if task_name and task_version:
+            task_id = task_definition_pb2.TaskIdentifier(
+                name=task_name,
+                project=cfg.project,
+                domain=cfg.domain,
+                org=cfg.org,
+                version=task_version,
+            )
+        elif task_name:
+            task_name_id = task_definition_pb2.TaskName(
+                name=task_name,
+                project=cfg.project,
+                domain=cfg.domain,
+                org=cfg.org,
+            )
+        else:
+            project_id = identifier_pb2.ProjectIdentifier(
+                organization=cfg.org,
+                domain=cfg.domain,
+                name=cfg.project,
+            )
+
+        while True:
+            resp = await get_client().trigger_service.ListTriggers(
+                request=trigger_service_pb2.ListTriggersRequest(
+                    project_id=project_id,
+                    task_id=task_id,
+                    task_name=task_name_id,
+                    request=list_pb2.ListRequest(
+                        limit=limit,
+                        token=token,
+                    ),
+                )
+            )
+            token = resp.token
+            for r in resp.triggers:
+                yield cls(r)
+            if not token:
+                break
+
+    @syncify
+    @classmethod
+    async def update(cls, name: str, task_name: str, active: bool):
+        """
+        Pause a trigger by its name and associated task name.
+        """
+        ensure_client()
+        cfg = get_init_config()
+        await get_client().trigger_service.UpdateTriggers(
+            request=trigger_service_pb2.UpdateTriggersRequest(
+                names=[
+                    identifier_pb2.TriggerName(
+                        org=cfg.org,
+                        project=cfg.project,
+                        domain=cfg.domain,
+                        name=name,
+                        task_name=task_name,
+                    )
+                ],
+                active=active,
+            )
+        )
+
+    @syncify
+    @classmethod
+    async def delete(cls, name: str, task_name: str):
+        """
+        Delete a trigger by its name.
+        """
+        ensure_client()
+        cfg = get_init_config()
+        await get_client().trigger_service.DeleteTriggers(
+            request=trigger_service_pb2.DeleteTriggersRequest(
+                names=[
+                    identifier_pb2.TriggerName(
+                        org=cfg.org,
+                        project=cfg.project,
+                        domain=cfg.domain,
+                        name=name,
+                        task_name=task_name,
+                    )
+                ],
+            )
+        )
+
+    @property
+    def id(self) -> identifier_pb2.TriggerIdentifier:
+        return self.pb2.id
+
+    @property
+    def name(self) -> str:
+        return self.id.name.name
+
+    @property
+    def task_name(self) -> str:
+        return self.id.name.task_name
+
+    @property
+    def automation_spec(self) -> common_pb2.TriggerAutomationSpec:
+        return self.pb2.automation_spec
+
+    async def get_details(self) -> TriggerDetails:
+        """
+        Get detailed information about this trigger.
+        """
+        if not self.details:
+            details = await TriggerDetails.get.aio(name=self.pb2.id.name.name)
+            self.details = details
+        return self.details
+
+    @property
+    def is_active(self) -> bool:
+        return self.pb2.active
+
+    def _rich_automation(self, automation: common_pb2.TriggerAutomationSpec):
+        if automation.type == common_pb2.TriggerAutomationSpec.type.TYPE_NONE:
+            yield "none", None
+        elif automation.type == common_pb2.TriggerAutomationSpec.type.TYPE_SCHEDULE:
+            if automation.schedule.cron is not None:
+                yield "cron", automation.schedule.cron
+            elif automation.schedule.rate is not None:
+                r = automation.schedule.rate
+                yield (
+                    "fixed_rate",
+                    (
+                        f"Every [{r.value}] {r.unit} starting at "
+                        f"{r.start_time.ToDatetime() if automation.HasField('start_time') else 'now'}"
+                    ),
+                )
+
+    def __rich_repr__(self):
+        yield "task_name", self.task_name
+        yield "name", self.name
+        yield from self._rich_automation(self.pb2.automation_spec)
+        yield "auto_activate", self.is_active

flyte/remote/_user.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from flyteidl.service import identity_pb2
+from flyteidl.service.identity_pb2 import UserInfoResponse
+
+from .._initialize import ensure_client, get_client
+from ..syncify import syncify
+from ._common import ToJSONMixin
+
+
+@dataclass
+class User(ToJSONMixin):
+    pb2: UserInfoResponse
+
+    @syncify
+    @classmethod
+    async def get(cls) -> User:
+        """
+        Fetches information about the currently logged in user.
+        Returns: A User object containing details about the user.
+        """
+        ensure_client()
+
+        resp = await get_client().identity_service.UserInfo(identity_pb2.UserInfoRequest())
+        return cls(resp)
+
+    def subject(self) -> str:
+        return self.pb2.subject
+
+    def name(self) -> str:
+        return self.pb2.name

flyte/report/_report.py

@@ -4,7 +4,6 @@ import string
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Dict, List, Union
 
-from flyte._internal.runtime import io
 from flyte._logging import logger
 from flyte._tools import ipython_check
 from flyte.syncify import syncify
@@ -133,6 +132,7 @@ async def flush():
     """
     import flyte.storage as storage
     from flyte._context import internal_ctx
+    from flyte._internal.runtime import io
 
     if not internal_ctx().is_task_context():
         return

flyte/storage/__init__.py

@@ -3,6 +3,8 @@ __all__ = [
     "GCS",
     "S3",
     "Storage",
+    "exists",
+    "exists_sync",
     "get",
     "get_configured_fsspec_kwargs",
     "get_random_local_directory",
@@ -11,13 +13,15 @@ __all__ = [
     "get_underlying_filesystem",
     "is_remote",
     "join",
+    "open",
     "put",
     "put_stream",
-    "put_stream",
 ]
 
 from ._config import ABFS, GCS, S3, Storage
 from ._storage import (
+    exists,
+    exists_sync,
     get,
     get_configured_fsspec_kwargs,
     get_random_local_directory,
@@ -26,6 +30,7 @@ from ._storage import (
     get_underlying_filesystem,
     is_remote,
     join,
+    open,
     put,
     put_stream,
 )

flyte/storage/_config.py

@@ -61,6 +61,7 @@ class S3(Storage):
     endpoint: typing.Optional[str] = None
     access_key_id: typing.Optional[str] = None
     secret_access_key: typing.Optional[str] = None
+    region: typing.Optional[str] = None
 
     _KEY_ENV_VAR_MAPPING: ClassVar[typing.Dict[str, str]] = {
         "endpoint": "FLYTE_AWS_ENDPOINT",
@@ -76,7 +77,7 @@ class S3(Storage):
     _KEY_SKIP_SIGNATURE: ClassVar = "skip_signature"
 
     @classmethod
-    def auto(cls) -> S3:
+    def auto(cls, region: str | None = None) -> S3:
         """
         :return: Config
         """
@@ -88,6 +89,7 @@ class S3(Storage):
         kwargs = set_if_exists(kwargs, "endpoint", endpoint)
         kwargs = set_if_exists(kwargs, "access_key_id", access_key_id)
         kwargs = set_if_exists(kwargs, "secret_access_key", secret_access_key)
+        kwargs = set_if_exists(kwargs, "region", region)
 
         return S3(**kwargs)
 
@@ -141,6 +143,8 @@ class S3(Storage):
             kwargs["config"] = config
         kwargs["client_options"] = client_options or None
         kwargs["retry_config"] = retry_config or None
+        if self.region:
+            kwargs["region"] = self.region
 
         return kwargs