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

flyte/remote/_secret.py

@@ -4,9 +4,9 @@ from dataclasses import dataclass
 from typing import AsyncIterator, Literal, Union
 
 import rich.repr
+from flyteidl2.secret import definition_pb2, payload_pb2
 
-from flyte._initialize import ensure_client, get_client, get_common_config
-from flyte._protos.secret import definition_pb2, payload_pb2
+from flyte._initialize import ensure_client, get_client, get_init_config
 from flyte.remote._common import ToJSONMixin
 from flyte.syncify import syncify
 
@@ -21,12 +21,19 @@ class Secret(ToJSONMixin):
     @classmethod
     async def create(cls, name: str, value: Union[str, bytes], type: SecretTypes = "regular"):
         ensure_client()
-        cfg = get_common_config()
-        secret_type = (
-            definition_pb2.SecretType.SECRET_TYPE_GENERIC
-            if type == "regular"
-            else definition_pb2.SecretType.SECRET_TYPE_IMAGE_PULL_SECRET
-        )
+        cfg = get_init_config()
+        project = cfg.project
+        domain = cfg.domain
+
+        if type == "regular":
+            secret_type = definition_pb2.SecretType.SECRET_TYPE_GENERIC
+
+        else:
+            secret_type = definition_pb2.SecretType.SECRET_TYPE_IMAGE_PULL_SECRET
+            if project or domain:
+                raise ValueError(
+                    f"Project `{project}` or domain `{domain}` should not be set when creating the image pull secret."
+                )
 
         if isinstance(value, str):
             secret = definition_pb2.SecretSpec(
@@ -42,8 +49,8 @@ class Secret(ToJSONMixin):
             request=payload_pb2.CreateSecretRequest(
                 id=definition_pb2.SecretIdentifier(
                     organization=cfg.org,
-                    project=cfg.project,
-                    domain=cfg.domain,
+                    project=project,
+                    domain=domain,
                     name=name,
                 ),
                 secret_spec=secret,
@@ -54,7 +61,7 @@ class Secret(ToJSONMixin):
     @classmethod
     async def get(cls, name: str) -> Secret:
         ensure_client()
-        cfg = get_common_config()
+        cfg = get_init_config()
         resp = await get_client().secrets_service.GetSecret(
             request=payload_pb2.GetSecretRequest(
                 id=definition_pb2.SecretIdentifier(
@@ -71,29 +78,31 @@ class Secret(ToJSONMixin):
     @classmethod
     async def listall(cls, limit: int = 100) -> AsyncIterator[Secret]:
         ensure_client()
-        cfg = get_common_config()
-        token = None
+        cfg = get_init_config()
+        per_cluster_tokens = None
         while True:
             resp = await get_client().secrets_service.ListSecrets(  # type: ignore
                 request=payload_pb2.ListSecretsRequest(
                     organization=cfg.org,
                     project=cfg.project,
                     domain=cfg.domain,
-                    token=token,
+                    per_cluster_tokens=per_cluster_tokens,
                     limit=limit,
                 ),
             )
-            token = resp.token
+            per_cluster_tokens = resp.per_cluster_tokens
+            round_items = [v for _, v in per_cluster_tokens.items() if v]
+            has_next = any(round_items)
             for r in resp.secrets:
                 yield cls(r)
-            if not token:
+            if not has_next:
                 break
 
     @syncify
     @classmethod
     async def delete(cls, name):
         ensure_client()
-        cfg = get_common_config()
+        cfg = get_init_config()
         await get_client().secrets_service.DeleteSecret(  # type: ignore
             request=payload_pb2.DeleteSecretRequest(
                 id=definition_pb2.SecretIdentifier(

flyte/remote/_task.py

@@ -6,19 +6,18 @@ from dataclasses import dataclass
 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
 
 
@@ -99,6 +98,7 @@ AutoVersioning = Literal["latest", "current"]
 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(
@@ -150,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,
@@ -284,6 +284,13 @@ class TaskDetails(ToJSONMixin):
             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,
         *,
@@ -295,6 +302,7 @@ class TaskDetails(ToJSONMixin):
         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:
@@ -342,7 +350,11 @@ class TaskDetails(ToJSONMixin):
             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)
+        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:
         """
@@ -438,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/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,
 )