flyte 0.0.1b0__py3-none-any.whl

union/_internal/controllers/__init__.py

@@ -0,0 +1,77 @@
+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

@@ -0,0 +1,77 @@
+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

@@ -0,0 +1,39 @@
+# 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

@@ -0,0 +1,40 @@
+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

@@ -0,0 +1,131 @@
+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

@@ -0,0 +1,43 @@
+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

@@ -0,0 +1,169 @@
+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)