flyte 0.1.0__py3-none-any.whl → 0.2.0a0__py3-none-any.whl

flyte/models.py

@@ -0,0 +1,391 @@
+from __future__ import annotations
+
+import inspect
+import os
+import pathlib
+import tempfile
+from dataclasses import dataclass, field, replace
+from typing import TYPE_CHECKING, Any, Callable, ClassVar, Dict, Literal, Optional, Tuple, Type
+
+import rich.repr
+
+from flyte._docstring import Docstring
+from flyte._interface import extract_return_annotation
+from flyte._logging import logger
+from flyte._utils.helpers import base36_encode
+
+if TYPE_CHECKING:
+    from flyteidl.core import literals_pb2
+
+    from flyte._internal.imagebuild.image_builder import ImageCache
+    from flyte.report import Report
+
+
+def generate_random_name() -> str:
+    """
+    Generate a random name for the task. This is used to create unique names for tasks.
+    TODO we can use unique-namer in the future, for now its just guids
+    """
+    from uuid import uuid4
+
+    return str(uuid4())  # Placeholder for actual random name generation logic
+
+
+@rich.repr.auto
+@dataclass(frozen=True, kw_only=True)
+class ActionID:
+    """
+    A class representing the ID of an Action, nested within a Run. This is used to identify a specific action on a task.
+    """
+
+    name: str
+    run_name: str | None = None
+    project: str | None = None
+    domain: str | None = None
+    org: str | None = None
+
+    def __post_init__(self):
+        if self.run_name is None:
+            object.__setattr__(self, "run_name", self.name)
+
+    @classmethod
+    def create_random(cls):
+        name = generate_random_name()
+        return cls(name=name, run_name=name)
+
+    def new_sub_action(self, name: str | None = None) -> ActionID:
+        """
+        Create a new sub-run with the given name. If  name is None, a random name will be generated.
+        """
+        if name is None:
+            name = generate_random_name()
+        return replace(self, name=name)
+
+    def new_sub_action_from(self, task_call_seq: int, task_hash: str, input_hash: str, group: str | None) -> ActionID:
+        """Make a deterministic name"""
+        import hashlib
+
+        components = f"{self.name}-{input_hash}-{task_hash}-{task_call_seq}" + (f"-{group}" if group else "")
+        logger.debug(f"----- Generating sub-action ID from components: {components}")
+        # has the components into something deterministic
+        bytes_digest = hashlib.md5(components.encode()).digest()
+        new_name = base36_encode(bytes_digest)
+        return self.new_sub_action(new_name)
+
+
+@rich.repr.auto
+@dataclass(frozen=True, kw_only=True)
+class RawDataPath:
+    """
+    A class representing the raw data path for a task. This is used to store the raw data for the task execution and
+    also get mutations on the path.
+    """
+
+    path: str
+
+    @classmethod
+    def from_local_folder(cls, local_folder: str | pathlib.Path | None = None) -> RawDataPath:
+        """
+        Create a new context attribute object, with local path given. Will be created if it doesn't exist.
+        :return: Path to the temporary directory
+        """
+        match local_folder:
+            case pathlib.Path():
+                local_folder.mkdir(parents=True, exist_ok=True)
+                return RawDataPath(path=str(local_folder))
+            case None:
+                # Create a temporary directory for data storage
+                p = tempfile.mkdtemp()
+                logger.debug(f"Creating temporary directory for data storage: {p}")
+                pathlib.Path(p).mkdir(parents=True, exist_ok=True)
+                return RawDataPath(path=p)
+            case str():
+                return RawDataPath(path=local_folder)
+            case _:
+                raise ValueError(f"Invalid local path {local_folder}")
+
+    def get_random_remote_path(self, file_name: Optional[str] = None) -> str:
+        """
+        Returns a random path for uploading a file/directory to.
+
+        :param file_name: If given, will be joined after a randomly generated portion.
+        :return:
+        """
+        import random
+        from uuid import UUID
+
+        import fsspec
+        from fsspec.utils import get_protocol
+
+        random_string = UUID(int=random.getrandbits(128)).hex
+        file_prefix = self.path
+
+        protocol = get_protocol(file_prefix)
+        if "file" in protocol:
+            local_path = pathlib.Path(file_prefix) / random_string
+            if file_name:
+                # Only if file name is given do we create the parent, because it may be needed as a folder otherwise
+                local_path = local_path / file_name
+                if not local_path.exists():
+                    local_path.parent.mkdir(exist_ok=True, parents=True)
+                    local_path.touch()
+            return str(local_path.absolute())
+
+        fs = fsspec.filesystem(protocol)
+        if file_prefix.endswith(fs.sep):
+            file_prefix = file_prefix[:-1]
+        remote_path = fs.sep.join([file_prefix, random_string])
+        if file_name:
+            remote_path = fs.sep.join([remote_path, file_name])
+        return remote_path
+
+
+@rich.repr.auto
+@dataclass(frozen=True)
+class GroupData:
+    name: str
+
+
+@rich.repr.auto
+@dataclass(frozen=True, kw_only=True)
+class TaskContext:
+    """
+    A context class to hold the current task executions context.
+    This can be used to access various contextual parameters in the task execution by the user.
+
+    :param action: The action ID of the current execution. This is always set, within a run.
+    :param version: The version of the executed task. This is set when the task is executed by an action and will be
+      set on all sub-actions.
+    """
+
+    action: ActionID
+    version: str
+    raw_data_path: RawDataPath
+    output_path: str
+    run_base_dir: str
+    report: Report
+    group_data: GroupData | None = None
+    checkpoints: Checkpoints | None = None
+    code_bundle: CodeBundle | None = None
+    compiled_image_cache: ImageCache | None = None
+    data: Dict[str, Any] = field(default_factory=dict)
+    mode: Literal["local", "remote", "hybrid"] = "remote"
+
+    def replace(self, **kwargs) -> TaskContext:
+        if "data" in kwargs:
+            rec_data = kwargs.pop("data")
+            if rec_data is None:
+                return replace(self, **kwargs)
+            data = {}
+            if self.data is not None:
+                data = self.data.copy()
+            data.update(rec_data)
+            kwargs.update({"data": data})
+        return replace(self, **kwargs)
+
+    def __getitem__(self, key: str) -> Optional[Any]:
+        return self.data.get(key)
+
+
+@rich.repr.auto
+@dataclass(frozen=True, kw_only=True)
+class CodeBundle:
+    """
+    A class representing a code bundle for a task. This is used to package the code and the inflation path.
+    The code bundle computes the version of the code using the hash of the code.
+
+    :param computed_version: The version of the code bundle. This is the hash of the code.
+    :param destination: The destination path for the code bundle to be inflated to.
+    :param tgz: Optional path to the tgz file.
+    :param pkl: Optional path to the pkl file.
+    :param downloaded_path: The path to the downloaded code bundle. This is only available during runtime, when
+        the code bundle has been downloaded and inflated.
+    """
+
+    computed_version: str
+    destination: str = "."
+    tgz: str | None = None
+    pkl: str | None = None
+    downloaded_path: pathlib.Path | None = None
+
+    # runtime_dependencies: Tuple[str, ...] = field(default_factory=tuple)  In the future if we want we could add this
+    # but this messes up actors, spark etc
+
+    def __post_init__(self):
+        if self.tgz is None and self.pkl is None:
+            raise ValueError("Either tgz or pkl must be provided")
+
+    def with_downloaded_path(self, path: pathlib.Path) -> CodeBundle:
+        """
+        Create a new CodeBundle with the given downloaded path.
+        """
+        return replace(self, downloaded_path=path)
+
+
+@rich.repr.auto
+@dataclass(frozen=True)
+class Checkpoints:
+    """
+    A class representing the checkpoints for a task. This is used to store the checkpoints for the task execution.
+    """
+
+    prev_checkpoint_path: str | None
+    checkpoint_path: str | None
+
+
+class _has_default:
+    """
+    A marker class to indicate that a specific input has a default value or not.
+    This is used to determine if the input is required or not.
+    """
+
+
+@dataclass(frozen=True)
+class NativeInterface:
+    """
+    A class representing the native interface for a task. This is used to interact with the task and its execution
+    context.
+    """
+
+    inputs: Dict[str, Tuple[Type, Any]]
+    outputs: Dict[str, Type]
+    docstring: Optional[Docstring] = None
+
+    # This field is used to indicate that the task has a default value for the input, but already in the
+    # remote form.
+    _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:
+        """
+        Check if the task has outputs. This is used to determine if the task has outputs or not.
+        """
+        return self.outputs is not None and len(self.outputs) > 0
+
+    def num_required_inputs(self) -> int:
+        """
+        Get the number of required inputs for the task. This is used to determine how many inputs are required for the
+        task execution.
+        """
+        return sum(1 for t in self.inputs.values() if t[1] is inspect.Parameter.empty)
+
+    @classmethod
+    def from_types(
+        cls,
+        inputs: Dict[str, Tuple[Type, Type[_has_default] | Type[inspect._empty]]],
+        outputs: Dict[str, Type],
+        default_inputs: Optional[Dict[str, literals_pb2.Literal]] = None,
+    ) -> NativeInterface:
+        """
+        Create a new NativeInterface from the given types. This is used to create a native interface for the task.
+        :param inputs: A dictionary of input names and their types and a value indicating if they have a default value.
+        :param outputs: A dictionary of output names and their types.
+        :param default_inputs: Optional dictionary of default inputs for remote tasks.
+        :return: A NativeInterface object with the given inputs and outputs.
+        """
+        for k, v in inputs.items():
+            if v[1] is cls.has_default and (default_inputs is None or k not in default_inputs):
+                raise ValueError(f"Input {k} has a default value but no default input provided for remote task.")
+        return cls(inputs=inputs, outputs=outputs, _remote_defaults=default_inputs)
+
+    @classmethod
+    def from_callable(cls, func: Callable) -> NativeInterface:
+        """
+        Extract the native interface from the given function. This is used to create a native interface for the task.
+        """
+        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()}
+
+        # Get return type
+        outputs = extract_return_annotation(sig.return_annotation)
+        return cls(inputs=param_info, outputs=outputs)
+
+    def convert_to_kwargs(self, *args, **kwargs) -> Dict[str, Any]:
+        """
+        Convert the given arguments to keyword arguments based on the native interface. This is used to convert the
+        arguments to the correct types for the task execution.
+        """
+        # Convert positional arguments to keyword arguments
+        if len(args) > len(self.inputs):
+            raise ValueError(f"Too many positional arguments provided, inputs {self.inputs.keys()}, args {len(args)}")
+        for arg, input_name in zip(args, self.inputs.keys()):
+            kwargs[input_name] = arg
+        return kwargs
+
+    def get_input_types(self) -> Dict[str, Type]:
+        """
+        Get the input types for the task. This is used to get the types of the inputs for the task execution.
+        """
+        return {k: v[0] for k, v in self.inputs.items()}
+
+    def __repr__(self):
+        """
+        Returns a string representation of the task interface.
+        """
+        i = "("
+        if self.inputs:
+            initial = True
+            for key, tpe in self.inputs.items():
+                if not initial:
+                    i += ", "
+                initial = False
+                tp = tpe[0] if isinstance(tpe[0], str) else tpe[0].__name__
+                i += f"{key}: {tp}"
+                if tpe[1] is not inspect.Parameter.empty:
+                    if tpe[1] is self.has_default:
+                        i += " = ..."
+                    else:
+                        i += f" = {tpe[1]}"
+        i += ")"
+        if self.outputs:
+            initial = True
+            multi = len(self.outputs) > 1
+            i += " -> "
+            if multi:
+                i += "("
+            for key, tpe in self.outputs.items():
+                if not initial:
+                    i += ", "
+                initial = False
+                tp = tpe.__name__ if isinstance(tpe, type) else tpe
+                i += f"{key}: {tp}"
+            if multi:
+                i += ")"
+        return i + ":"
+
+
+@dataclass
+class SerializationContext:
+    """
+    This object holds serialization time contextual information, that can be used when serializing the task and
+    various parameters of a tasktemplate. This is only available when the task is being serialized and can be
+    during a deployment or runtime.
+
+    :param version: The version of the task
+    :param code_bundle: The code bundle for the task. This is used to package the code and the inflation path.
+    :param input_path: The path to the inputs for the task. This is used to determine where the inputs will be located
+    :param output_path: The path to the outputs for the task. This is used to determine where the outputs will be
+     located
+    """
+
+    version: str
+    project: str | None = None
+    domain: str | None = None
+    org: str | None = None
+    code_bundle: Optional[CodeBundle] = None
+    input_path: str = "{{.input}}"
+    output_path: str = "{{.outputPrefix}}"
+    _entrypoint_path: str = field(default="_bin/runtime.py", init=False)
+    image_cache: ImageCache | None = None
+    root_dir: Optional[pathlib.Path] = None
+
+    def get_entrypoint_path(self, interpreter_path: str) -> str:
+        """
+        Get the entrypoint path for the task. This is used to determine the entrypoint for the task execution.
+        :param interpreter_path: The path to the interpreter (python)
+        """
+        return os.path.join(os.path.dirname(os.path.dirname(interpreter_path)), self._entrypoint_path)

flyte/remote/__init__.py

@@ -0,0 +1,26 @@
+"""
+Remote Entities that are accessible from the Union Server once deployed or created.
+"""
+
+__all__ = [
+    "Action",
+    "ActionDetails",
+    "ActionInputs",
+    "ActionOutputs",
+    "Project",
+    "Run",
+    "RunDetails",
+    "Secret",
+    "SecretTypes",
+    "Task",
+    "create_channel",
+    "upload_dir",
+    "upload_file",
+]
+
+from ._client.auth import create_channel
+from ._data import upload_dir, upload_file
+from ._project import Project
+from ._run import Action, ActionDetails, ActionInputs, ActionOutputs, Run, RunDetails
+from ._secret import Secret, SecretTypes
+from ._task import Task

flyte/remote/_client/_protocols.py

@@ -0,0 +1,133 @@
+from typing import AsyncIterator, Protocol
+
+from flyteidl.admin import project_attributes_pb2, project_pb2, version_pb2
+from flyteidl.service import dataproxy_pb2
+from grpc.aio import UnaryStreamCall
+from grpc.aio._typing import RequestType
+
+from flyte._protos.secret import payload_pb2
+from flyte._protos.workflow import run_logs_service_pb2, run_service_pb2, task_service_pb2
+
+
+class MetadataServiceProtocol(Protocol):
+    async def GetVersion(self, request: version_pb2.GetVersionRequest) -> version_pb2.GetVersionResponse: ...
+
+
+class ProjectDomainService(Protocol):
+    async def RegisterProject(
+        self, request: project_pb2.ProjectRegisterRequest
+    ) -> project_pb2.ProjectRegisterResponse: ...
+
+    async def UpdateProject(self, request: project_pb2.Project) -> project_pb2.ProjectUpdateResponse: ...
+
+    async def GetProject(self, request: project_pb2.ProjectGetRequest) -> project_pb2.Project: ...
+
+    async def ListProjects(self, request: project_pb2.ProjectListRequest) -> project_pb2.Projects: ...
+
+    async def GetDomains(self, request: project_pb2.GetDomainRequest) -> project_pb2.GetDomainsResponse: ...
+
+    async def UpdateProjectDomainAttributes(
+        self, request: project_attributes_pb2.ProjectAttributesUpdateRequest
+    ) -> project_pb2.ProjectUpdateResponse: ...
+
+    async def GetProjectDomainAttributes(
+        self, request: project_attributes_pb2.ProjectAttributesGetRequest
+    ) -> project_attributes_pb2.ProjectAttributes: ...
+
+    async def DeleteProjectDomainAttributes(
+        self, request: project_attributes_pb2.ProjectAttributesDeleteRequest
+    ) -> project_attributes_pb2.ProjectAttributesDeleteResponse: ...
+
+    async def UpdateProjectAttributes(
+        self, request: project_attributes_pb2.ProjectAttributesUpdateRequest
+    ) -> project_attributes_pb2.ProjectAttributesUpdateResponse: ...
+
+    async def GetProjectAttributes(
+        self, request: project_attributes_pb2.ProjectAttributesGetRequest
+    ) -> project_attributes_pb2.ProjectAttributes: ...
+
+    async def DeleteProjectAttributes(
+        self, request: project_attributes_pb2.ProjectAttributesDeleteRequest
+    ) -> project_attributes_pb2.ProjectAttributesDeleteResponse: ...
+
+
+class TaskService(Protocol):
+    async def DeployTask(self, request: task_service_pb2.DeployTaskRequest) -> task_service_pb2.DeployTaskResponse: ...
+
+    async def GetTaskDetails(
+        self, request: task_service_pb2.GetTaskDetailsRequest
+    ) -> task_service_pb2.GetTaskDetailsResponse: ...
+
+    async def ListTasks(self, request: task_service_pb2.ListTasksRequest) -> task_service_pb2.ListTasksResponse: ...
+
+
+class RunService(Protocol):
+    async def CreateRun(self, request: run_service_pb2.CreateRunRequest) -> run_service_pb2.CreateRunResponse: ...
+
+    async def AbortRun(self, request: run_service_pb2.AbortRunRequest) -> run_service_pb2.AbortRunResponse: ...
+
+    async def GetRunDetails(
+        self, request: run_service_pb2.GetRunDetailsRequest
+    ) -> run_service_pb2.GetRunDetailsResponse: ...
+
+    async def WatchRunDetails(
+        self, request: run_service_pb2.WatchRunDetailsRequest
+    ) -> AsyncIterator[run_service_pb2.WatchRunDetailsResponse]: ...
+
+    async def GetActionDetails(
+        self, request: run_service_pb2.GetActionDetailsRequest
+    ) -> run_service_pb2.GetActionDetailsResponse: ...
+
+    async def WatchActionDetails(
+        self, request: run_service_pb2.WatchActionDetailsRequest
+    ) -> AsyncIterator[run_service_pb2.WatchActionDetailsResponse]: ...
+
+    async def GetActionData(
+        self, request: run_service_pb2.GetActionDataRequest
+    ) -> run_service_pb2.GetActionDataResponse: ...
+
+    async def ListRuns(self, request: run_service_pb2.ListRunsRequest) -> run_service_pb2.ListRunsResponse: ...
+
+    async def WatchRuns(
+        self, request: run_service_pb2.WatchRunsRequest
+    ) -> AsyncIterator[run_service_pb2.WatchRunsResponse]: ...
+
+    async def ListActions(self, request: run_service_pb2.ListActionsRequest) -> run_service_pb2.ListActionsResponse: ...
+
+    async def WatchActions(
+        self, request: run_service_pb2.WatchActionsRequest
+    ) -> AsyncIterator[run_service_pb2.WatchActionsResponse]: ...
+
+
+class DataProxyService(Protocol):
+    async def CreateUploadLocation(
+        self, request: dataproxy_pb2.CreateUploadLocationRequest
+    ) -> dataproxy_pb2.CreateUploadLocationResponse: ...
+
+    async def CreateDownloadLocation(
+        self, request: dataproxy_pb2.CreateDownloadLocationRequest
+    ) -> dataproxy_pb2.CreateDownloadLocationResponse: ...
+
+    async def CreateDownloadLink(
+        self, request: dataproxy_pb2.CreateDownloadLinkRequest
+    ) -> dataproxy_pb2.CreateDownloadLinkResponse: ...
+
+    async def GetData(self, request: dataproxy_pb2.GetDataRequest) -> dataproxy_pb2.GetDataResponse: ...
+
+
+class RunLogsService(Protocol):
+    def TailLogs(
+        self, request: run_logs_service_pb2.TailLogsRequest
+    ) -> UnaryStreamCall[RequestType, run_logs_service_pb2.TailLogsResponse]: ...
+
+
+class SecretService(Protocol):
+    async def CreateSecret(self, request: payload_pb2.CreateSecretRequest) -> payload_pb2.CreateSecretResponse: ...
+
+    async def UpdateSecret(self, request: payload_pb2.UpdateSecretRequest) -> payload_pb2.UpdateSecretResponse: ...
+
+    async def GetSecret(self, request: payload_pb2.GetSecretRequest) -> payload_pb2.GetSecretResponse: ...
+
+    async def ListSecrets(self, request: payload_pb2.ListSecretsRequest) -> payload_pb2.ListSecretsResponse: ...
+
+    async def DeleteSecret(self, request: payload_pb2.DeleteSecretRequest) -> payload_pb2.DeleteSecretResponse: ...

flyte/remote/_client/auth/__init__.py

@@ -0,0 +1,12 @@
+from flyte.remote._client.auth._channel import create_channel
+from flyte.remote._client.auth._client_config import AuthType, ClientConfig
+from flyte.remote._client.auth.errors import AccessTokenNotFoundError, AuthenticationError, AuthenticationPending
+
+__all__ = [
+    "AccessTokenNotFoundError",
+    "AuthType",
+    "AuthenticationError",
+    "AuthenticationPending",
+    "ClientConfig",
+    "create_channel",
+]

flyte/remote/_client/auth/_auth_utils.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+import base64
+from typing import Literal
+
+
+def decode_api_key(encoded_str: str) -> tuple[str, str, str, str | Literal["None"]]:
+    """Decode encoded base64 string into app credentials. endpoint, client_id, client_secret, org"""
+    endpoint, client_id, client_secret, org = base64.b64decode(encoded_str.encode("utf-8")).decode("utf-8").split(":")
+    # For consistency, let's make sure org is always a non-empty string
+    if not org:
+        org = "None"
+
+    return endpoint, client_id, client_secret, org