flyte 0.0.1b0__py3-none-any.whl → 2.0.0b46__py3-none-any.whl

union/_internal/controllers/__init__.py

@@ -1,77 +0,0 @@
-import threading
-from typing import Any, Literal, Optional, Protocol
-
-from union._datastructures import ActionID
-from union._task import TaskTemplate
-
-__all__ = ["Controller", "ControllerType", "create_controller", "get_controller"]
-
-ControllerType = Literal["local", "remote"]
-
-
-class Controller(Protocol):
-    """
-    Controller interface, that is used to execute tasks. The implementation of this interface,
-    can execute tasks in different ways, such as locally, remotely etc.
-    """
-
-    async def submit(self, _task: TaskTemplate, *args, **kwargs) -> Any:
-        """
-        Submit a node to the controller asynchronously and wait for the result. This is async and will block
-        the current coroutine until the result is available.
-        """
-        ...
-
-    async def finalize_parent_action(self, action: ActionID):
-        """
-        Finalize the parent action. This can be called to cleanup the action and should be called after the parent
-        task completes
-        :param action: Action ID
-        :return:
-        """
-        ...
-
-    def stop(self):
-        """
-        Stops the engine and should be called when the engine is no longer needed.
-        """
-        ...
-
-
-# Internal state holder
-class _ControllerState:
-    controller: Optional[Controller] = None
-    lock = threading.Lock()
-
-
-async def get_controller() -> Controller:
-    """
-    Get the controller instance. Raise an error if it has not been created.
-    """
-    if _ControllerState.controller is not None:
-        return _ControllerState.controller
-    raise RuntimeError("Controller is not initialized. Please call get_or_create_controller() first.")
-
-
-def create_controller(
-    ct: ControllerType,
-    **kwargs,
-) -> Controller:
-    """
-    Create a new instance of the controller, based on the kind and the given configuration.
-    """
-    match ct:
-        case "local":
-            from ._local_controller import LocalController
-
-            controller = LocalController()
-        case ("remote" | "hybrid"):
-            from union._internal.controllers.remote import create_remote_controller
-
-            controller = create_remote_controller(**kwargs)
-        case _:
-            raise ValueError(f"{ct} is not a valid controller type.")
-
-    with _ControllerState.lock:
-        _ControllerState.controller = controller
-        return controller

union/_internal/controllers/_local_controller.py

@@ -1,77 +0,0 @@
-from typing import Any, Tuple
-
-from union import storage
-from union._context import internal_ctx
-from union._datastructures import ActionID, RawDataPath
-from union._internal.controllers import pbhash
-from union._internal.runtime.convert import (
-    Inputs,
-    convert_error_to_native,
-    convert_from_native_to_inputs,
-    convert_outputs_to_native,
-)
-from union._internal.runtime.entrypoints import direct_dispatch
-from union._logging import log, logger
-from union._task import TaskTemplate
-
-
-class LocalController:
-    def __init__(self):
-        logger.debug("LocalController init")
-
-    def _get_run_params(self, inputs: Inputs) -> Tuple[ActionID, RawDataPath]:
-        ctx = internal_ctx()
-        parent_run = ctx.data.task_context.action
-        # TODO assuming the raw_data_path is local, and for now not getting manipulated by the controller
-        # We will need to change this in case of remote execution, or create data sandboxes.
-        new_raw_data_path = ctx.data.raw_data_path
-        # TODO ideally we should generate the name deterministically using the inputs etc
-        sub_run = parent_run.new_sub_action()
-        return sub_run, new_raw_data_path
-
-    @log
-    async def submit(self, _task: TaskTemplate, *args, **kwargs) -> Any:
-        """
-        Submit a node to the controller
-        """
-        ctx = internal_ctx()
-        tctx = ctx.data.task_context
-        current_action_id = tctx.action
-        current_output_path = tctx.output_path
-
-        inputs = await convert_from_native_to_inputs(_task.native_interface, *args, **kwargs)
-        inputs_hash = self._input_hash(inputs)
-        sub_action_id = current_action_id.new_sub_action_from(
-            input_hash=inputs_hash,
-            group=tctx.group_data.name if tctx.group_data else None,
-        )
-        sub_action_output_path = storage.join(current_output_path, sub_action_id.name)
-        sub_action_raw_data_path = RawDataPath(path=sub_action_output_path)
-        out, err = await direct_dispatch(
-            _task,
-            controller=self,
-            action=sub_action_id,
-            raw_data_path=sub_action_raw_data_path,
-            inputs=inputs,
-            version=tctx.version,
-            checkpoints=tctx.checkpoints,
-            code_bundle=tctx.code_bundle,
-            output_path=sub_action_output_path,
-        )
-        if err:
-            raise convert_error_to_native(err)
-        if _task.native_interface.outputs:
-            out = await convert_outputs_to_native(_task.native_interface, out)
-        return out
-
-    async def finalize_parent_action(self, action: ActionID):
-        pass
-
-    def stop(self):
-        pass
-
-    def _input_hash(self, inputs: Inputs) -> str:
-        """
-        Returns the hash of the inputs
-        """
-        return pbhash.compute_hash_string(inputs.proto_inputs)

union/_internal/controllers/pbhash.py

@@ -1,39 +0,0 @@
-# This is a module that provides hashing utilities for Protobuf objects.
-import base64
-import hashlib
-import json
-
-from google.protobuf import json_format
-from google.protobuf.message import Message
-
-
-def compute_hash(pb: Message) -> bytes:
-    """
-    Computes a deterministic hash in bytes for the Protobuf object.
-    """
-    try:
-        pb_dict = json_format.MessageToDict(pb)
-        # json.dumps with sorted keys to ensure stability
-        stable_json_str = json.dumps(
-            pb_dict, sort_keys=True, separators=(",", ":")
-        )  # separators to ensure no extra spaces
-    except Exception as e:
-        raise ValueError(f"Failed to marshal Protobuf object {pb} to JSON with error: {e}")
-
-    try:
-        # Deterministically hash the JSON object to a byte array. Using SHA-256 for hashing here,
-        # assuming it provides a consistent hash output.
-        hash_obj = hashlib.sha256(stable_json_str.encode("utf-8"))
-    except Exception as e:
-        raise ValueError(f"Failed to hash JSON for Protobuf object {pb} with error: {e}")
-
-    # The digest is guaranteed to be 32 bytes long
-    return hash_obj.digest()
-
-
-def compute_hash_string(pb: Message) -> str:
-    """
-    Computes a deterministic hash in base64 encoded string for the Protobuf object
-    """
-    hash_bytes = compute_hash(pb)
-    return base64.b64encode(hash_bytes).decode("utf-8")

union/_internal/controllers/remote/__init__.py

@@ -1,40 +0,0 @@
-from typing import List
-
-from union.remote._client.auth import AuthType, ClientConfig
-
-from ._controller import RemoteController
-
-__all__ = ["RemoteController", "create_remote_controller"]
-
-
-def create_remote_controller(
-    *,
-    api_key: str | None = None,
-    auth_type: AuthType = "Pkce",
-    endpoint: str | None = None,
-    client_config: ClientConfig | None = None,
-    headless: bool = False,
-    insecure: bool = False,
-    insecure_skip_verify: bool = False,
-    ca_cert_file_path: str | None = None,
-    command: List[str] | None = None,
-    proxy_command: List[str] | None = None,
-    client_id: str | None = None,
-    client_credentials_secret: str | None = None,
-    rpc_retries: int = 3,
-    http_proxy_url: str | None = None,
-) -> RemoteController:
-    """
-    Create a new instance of the remote controller.
-    """
-    from ._client import ControllerClient
-    from ._controller import RemoteController
-
-    controller = RemoteController(
-        client_coro=ControllerClient.for_endpoint(
-            endpoint=endpoint, insecure=insecure, insecure_skip_verify=insecure_skip_verify
-        ),
-        workers=10,
-        max_system_retries=5,
-    )
-    return controller

union/_internal/controllers/remote/_action.py

@@ -1,131 +0,0 @@
-from __future__ import annotations
-
-from dataclasses import dataclass
-
-from flyteidl.core import execution_pb2
-
-from union._datastructures import GroupData
-from union._protos.workflow import run_definition_pb2, state_service_pb2, task_definition_pb2
-
-
-@dataclass
-class Action:
-    """
-    Coroutine safe, as we never do await operations in any method.
-    Holds the inmemory state of a task. It is combined representation of local and remote states.
-    """
-
-    action_id: run_definition_pb2.ActionIdentifier
-    parent_action_name: str
-    friendly_name: str | None = None
-    group: GroupData | None = None
-    task: task_definition_pb2.TaskSpec | None = None
-    inputs_uri: str | None = None
-    outputs_uri: str | None = None
-    err: execution_pb2.ExecutionError | None = None
-    phase: run_definition_pb2.Phase | None = None
-    started: bool = False
-    retries: int = 0
-    client_err: Exception | None = None  # This error is set when something goes wrong in the controller.
-
-    @property
-    def name(self) -> str:
-        return self.action_id.name
-
-    @property
-    def run_name(self) -> str:
-        return self.action_id.run.name
-
-    def is_terminal(self) -> bool:
-        """Check if resource has reached terminal state"""
-        if self.phase is None:
-            return False
-        return self.phase in [
-            run_definition_pb2.Phase.PHASE_FAILED,
-            run_definition_pb2.Phase.PHASE_SUCCEEDED,
-            run_definition_pb2.Phase.PHASE_ABORTED,
-            run_definition_pb2.Phase.PHASE_TIMED_OUT,
-        ]
-
-    def increment_retries(self):
-        self.retries += 1
-
-    def is_started(self) -> bool:
-        """Check if resource has been started."""
-        return self.started
-
-    def mark_started(self):
-        self.started = True
-        self.task = None
-
-    def merge_state(self, obj: state_service_pb2.ActionUpdate):
-        """
-        This method is invoked when the watch API sends an update about the state of the action. We need to merge
-        the state of the action with the current state of the action. It is possible that we have no phase information
-        prior to this.
-        :param obj:
-        :return:
-        """
-        if self.phase != obj.phase:
-            self.phase = obj.phase
-            self.err = obj.error if obj.HasField("error") else None
-        self.started = True
-
-    def merge_in_action_from_submit(self, action: Action):
-        """
-        This method is invoked when parent_action submits an action that was observed previously observed from the
-         watch. We need to merge in the contents of the action, while preserving the observed phase.
-
-        :param action: The submitted action
-        """
-        self.outputs_uri = action.outputs_uri
-        self.inputs_uri = action.inputs_uri
-        self.group = action.group
-        self.friendly_name = action.friendly_name
-        if not self.started:
-            self.task = action.task
-
-    def set_client_error(self, exc: Exception):
-        self.client_err = exc
-
-    def has_error(self) -> bool:
-        return self.client_err is not None or self.err is not None
-
-    @classmethod
-    def from_task(
-        cls,
-        parent_action_name: str,
-        sub_action_id: run_definition_pb2.ActionIdentifier,
-        group_data: GroupData,
-        task_spec: task_definition_pb2.TaskSpec,
-        inputs_uri: str,
-        outputs_prefix_uri: str,
-    ) -> Action:
-        return cls(
-            action_id=sub_action_id,
-            parent_action_name=parent_action_name,
-            friendly_name=task_spec.task_template.id.name,
-            group=group_data,
-            task=task_spec,
-            inputs_uri=inputs_uri,
-            outputs_uri=outputs_prefix_uri,
-        )
-
-    @classmethod
-    def from_state(cls, parent_action_name: str, obj: state_service_pb2.ActionUpdate) -> Action:
-        """
-        This creates a new action, from the watch api. This is possible in the case of a recovery, where the
-        stateservice knows about future actions and sends this information to the informer. We may not have encountered
-        the "task" itself yet, but we know about the action id and the state of the action.
-
-        :param parent_action_name:
-        :param obj:
-        :return:
-        """
-        return cls(
-            action_id=obj.action_id,
-            parent_action_name=parent_action_name,
-            phase=obj.phase,
-            started=True,
-            err=obj.error if obj.HasField("error") else None,
-        )

union/_internal/controllers/remote/_client.py

@@ -1,43 +0,0 @@
-from __future__ import annotations
-
-import grpc.aio
-
-from union._protos.workflow import queue_service_pb2_grpc, state_service_pb2_grpc
-from union.remote import create_channel
-
-from ._service_protocol import QueueService, StateService
-
-
-class ControllerClient:
-    """
-    A client for the Controller API.
-    """
-
-    def __init__(self, channel: grpc.aio.Channel):
-        self._channel = channel
-        self._state_service = state_service_pb2_grpc.StateServiceStub(channel=channel)
-        self._queue_service = queue_service_pb2_grpc.QueueServiceStub(channel=channel)
-
-    @classmethod
-    async def for_endpoint(cls, endpoint: str, *, insecure: bool = False, **kwargs) -> ControllerClient:
-        return cls(await create_channel(endpoint, insecure=insecure, **kwargs))
-
-    @property
-    def state_service(self) -> StateService:
-        """
-        The state service.
-        """
-        return self._state_service
-
-    @property
-    def queue_service(self) -> QueueService:
-        """
-        The queue service.
-        """
-        return self._queue_service
-
-    def close(self, grace: float | None = None):
-        """
-        Close the channel.
-        """
-        return self._channel.close(grace=grace)

union/_internal/controllers/remote/_controller.py

@@ -1,169 +0,0 @@
-from __future__ import annotations
-
-import asyncio
-from collections import defaultdict
-from typing import Any, Awaitable
-
-import union.storage as storage
-from union._code_bundle import build_pkl_bundle
-from union._context import internal_ctx
-from union._datastructures import ActionID, SerializationContext
-from union._internal.controllers import pbhash
-from union._internal.controllers.remote._action import Action
-from union._internal.controllers.remote._core import Controller
-from union._internal.controllers.remote._service_protocol import ClientSet
-from union._internal.runtime import io
-from union._internal.runtime.convert import (
-    Inputs,
-    convert_error_to_native,
-    convert_from_native_to_inputs,
-    convert_outputs_to_native,
-)
-from union._internal.runtime.task_serde import translate_task_to_wire
-from union._logging import logger
-from union._protos.workflow import run_definition_pb2
-from union._task import TaskTemplate
-from union.errors import RuntimeSystemError
-
-
-class RemoteController(Controller):
-    """
-    This a specialized controller that wraps the core controller and performs IO, serialization and deserialization
-    """
-
-    def __init__(
-        self,
-        client_coro: Awaitable[ClientSet],
-        workers: int,
-        max_system_retries: int,
-        default_parent_concurrency: int = 100,
-    ):
-        """ """
-        super().__init__(
-            client_coro=client_coro,
-            workers=workers,
-            max_system_retries=max_system_retries,
-        )
-        self._default_parent_concurrency = default_parent_concurrency
-        self._parent_action_semaphore = defaultdict(lambda: asyncio.Semaphore(default_parent_concurrency))
-
-    async def _submit(self, _task: TaskTemplate, *args, **kwargs) -> Any:
-        ctx = internal_ctx()
-        tctx = ctx.data.task_context
-        current_action_id = tctx.action
-        current_output_path = tctx.output_path
-
-        inputs = await convert_from_native_to_inputs(_task.native_interface, *args, **kwargs)
-        inputs_hash = self._input_hash(inputs)
-        sub_action_id = current_action_id.new_sub_action_from(
-            input_hash=inputs_hash,
-            group=tctx.group_data.name if tctx.group_data else None,
-        )
-        sub_run_output_path = storage.join(current_output_path, sub_action_id.name)
-        # In the case of a regular code bundle, we will just pass it down as it is to the downstream tasks
-        # It is not allowed to change the code bundle (for regular code bundles) in the middle of a run.
-        code_bundle = tctx.code_bundle
-
-        if code_bundle:
-            # but if we are using a pkl bundle, we need to build a new one for the downstream tasks
-            if code_bundle.pkl:
-                logger.debug(f"Building new pkl bundle for task {sub_action_id.name}")
-                code_bundle = await build_pkl_bundle(
-                    _task,
-                    upload_to_controlplane=False,
-                    upload_from_dataplane_path=io.pkl_path(sub_run_output_path),
-                )
-
-        inputs_uri = io.inputs_path(sub_run_output_path)
-        try:
-            # TODO Add retry decorator to this
-            await io.upload_inputs(inputs, inputs_uri)
-        except Exception as e:
-            logger.exception("Failed to upload inputs", e)
-            raise RuntimeSystemError(type(e).__name__, str(e)) from e
-        new_serialization_context = SerializationContext(
-            project=current_action_id.project,
-            domain=current_action_id.domain,
-            org=current_action_id.org,
-            code_bundle=code_bundle,
-            version=tctx.version,
-            # supplied version.
-            input_path=inputs_uri,
-            output_path=sub_run_output_path,
-            image_cache=ctx.data.task_context.compiled_image_cache,
-        )
-
-        task_spec = translate_task_to_wire(_task, new_serialization_context)
-
-        action = Action.from_task(
-            sub_action_id=run_definition_pb2.ActionIdentifier(
-                name=sub_action_id.name,
-                run=run_definition_pb2.RunIdentifier(
-                    name=current_action_id.run_name,
-                    project=current_action_id.project,
-                    domain=current_action_id.domain,
-                    org=current_action_id.org,
-                ),
-            ),
-            parent_action_name=current_action_id.name,
-            group_data=tctx.group_data,
-            task_spec=task_spec,
-            inputs_uri=inputs_uri,
-            outputs_prefix_uri=sub_run_output_path,
-        )
-
-        n = await self.submit_action(action)
-
-        if n.has_error() or n.phase == run_definition_pb2.PHASE_FAILED:
-            err = n.err or n.client_err
-            if not err and n.phase == run_definition_pb2.PHASE_FAILED:
-                logger.error(f"Server reported failure for action {n.action_id.name}, checking error file.")
-                error_path = io.error_path(n.outputs_uri)
-                # It is possible that the error file is not present in the case of a image pull failure or
-                # other reasons for failure. Ideally the err message should be sent by the server, but incase its
-                # missing, failed with unknown error
-                try:
-                    err = await io.load_error(error_path)
-                except Exception as e:
-                    logger.exception("Failed to load error file", e)
-                    err = RuntimeSystemError(
-                        type(e).__name__,
-                        f"Failed to load error file: {e}",
-                    )
-            else:
-                logger.error(f"Server reported failure for action {n.action_id.name}, error: {err}")
-            raise convert_error_to_native(err)
-
-        if _task.native_interface.outputs:
-            outputs_file_path = io.outputs_path(n.outputs_uri)
-            o = await io.load_outputs(outputs_file_path)
-            return await convert_outputs_to_native(_task.native_interface, o)
-        return None
-
-    async def submit(self, _task: TaskTemplate, *args, **kwargs) -> Any:
-        """
-        Submit a task to the remote controller.This creates a new action on the queue service.
-        """
-        from union._context import internal_ctx
-
-        ctx = internal_ctx()
-        current_action_id = ctx.data.task_context.action
-        async with self._parent_action_semaphore[current_action_id.name]:
-            return await self._submit(_task, *args, **kwargs)
-
-    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
-        to the control plane.
-        """
-        run_id = run_definition_pb2.RunIdentifier(
-            name=action_id.run_name,
-            project=action_id.project,
-            domain=action_id.domain,
-            org=action_id.org,
-        )
-        await super()._finalize_parent_action(run_id=run_id, parent_action_name=action_id.name)
-        self._parent_action_semaphore.pop(action_id.name, None)
-
-    def _input_hash(self, inputs: Inputs) -> str:
-        return pbhash.compute_hash_string(inputs.proto_inputs)