flyte 2.0.0b32__py3-none-any.whl

flyte/_internal/runtime/taskrunner.py

@@ -0,0 +1,209 @@
+"""
+This module is responsible for running tasks in the V2 runtime. All methods in this file should be
+invoked within a context tree.
+"""
+
+import pathlib
+import time
+from typing import Any, Dict, List, Optional, Tuple
+
+import flyte.report
+from flyte._context import internal_ctx
+from flyte._internal.imagebuild.image_builder import ImageCache
+from flyte._logging import log, logger
+from flyte._task import TaskTemplate
+from flyte.errors import CustomError, RuntimeSystemError, RuntimeUnknownError, RuntimeUserError
+from flyte.models import ActionID, Checkpoints, CodeBundle, RawDataPath, TaskContext
+
+from .. import Controller
+from .convert import (
+    Error,
+    Inputs,
+    Outputs,
+    convert_from_native_to_error,
+    convert_from_native_to_outputs,
+    convert_inputs_to_native,
+)
+from .io import load_inputs, upload_error, upload_outputs
+
+
+def replace_task_cli(args: List[str], inputs: Inputs, tmp_path: pathlib.Path, action: ActionID) -> List[str]:
+    """
+    This method can be used to run an task from the cli, if you have cli for the task. It will replace,
+    all the args with the task args.
+
+    The urun cli is of the format
+    ```python
+    ['urun', '--inputs', '{{.Inputs}}', '--outputs-path', '{{.Outputs}}', '--version', '',
+     '--raw-data-path', '{{.rawOutputDataPrefix}}',
+      '--checkpoint-path', '{{.checkpointOutputPrefix}}', '--prev-checkpoint', '{{.prevCheckpointPrefix}}',
+       '--run-name', '{{.runName}}', '--name', '{{.actionName}}',
+        '--tgz', 'some-path', '--dest', '.',
+         '--resolver', 'flyte._internal.resolvers.default.DefaultTaskResolver', '--resolver-args',
+          'mod', 'test_round_trip', 'instance', 'task1']
+    ```
+    We will replace, inputs, outputs, raw_data_path, checkpoint_path, prev_checkpoint, run_name, name
+    with supplied values.
+
+    :param args: urun command
+    :param inputs: converted inputs to the task
+    :param tmp_path: temporary path to use for the task
+    :param action: run id to use for the task
+    :return: modified args
+    """
+    # Iterate over all the args and replace the inputs, outputs, raw_data_path, checkpoint_path, prev_checkpoint,
+    # root_name, run_name with the supplied values
+    # first we will write the inputs to a file called inputs.pb
+    inputs_path = tmp_path / "inputs.pb"
+    with open(inputs_path, "wb") as f:
+        f.write(inputs.proto_inputs.SerializeToString())
+    # now modify the args
+    args = list(args)  # copy first because it's a proto container
+    for i, arg in enumerate(args):
+        match arg:
+            case "--inputs":
+                args[i + 1] = str(inputs_path)
+            case "--outputs-path":
+                args[i + 1] = str(tmp_path)
+            case "--raw-data-path":
+                args[i + 1] = str(tmp_path / "raw_data_path")
+            case "--checkpoint-path":
+                args[i + 1] = str(tmp_path / "checkpoint_path")
+            case "--prev-checkpoint":
+                args[i + 1] = str(tmp_path / "prev_checkpoint")
+            case "--run-name":
+                args[i + 1] = action.run_name or ""
+            case "--name":
+                args[i + 1] = action.name
+    insert_point = args.index("--raw-data-path")
+    args.insert(insert_point, str(tmp_path))
+    args.insert(insert_point, "--run-base-dir")
+    return args
+
+
+@log
+async def run_task(
+    tctx: TaskContext, controller: Controller, task: TaskTemplate, inputs: Dict[str, Any]
+) -> Tuple[Any, Optional[Exception]]:
+    try:
+        logger.info(f"Parent task executing {tctx.action}")
+        outputs = await task.execute(**inputs)
+        logger.info(f"Parent task completed successfully, {tctx.action}")
+        return outputs, None
+    except RuntimeSystemError as e:
+        logger.exception(f"Task failed with error: {e}")
+        return {}, e
+    except RuntimeUnknownError as e:
+        logger.exception(f"Task failed with error: {e}")
+        return {}, e
+    except RuntimeUserError as e:
+        logger.exception(f"Task failed with error: {e}")
+        return {}, e
+    except Exception as e:
+        logger.exception(f"Task failed with error: {e}")
+        return {}, CustomError.from_exception(e)
+    finally:
+        logger.info(f"Parent task finalized {tctx.action}")
+        # reconstruct run id here
+        await controller.finalize_parent_action(tctx.action)
+
+
+async def convert_and_run(
+    *,
+    task: TaskTemplate,
+    action: ActionID,
+    controller: Controller,
+    raw_data_path: RawDataPath,
+    version: str,
+    output_path: str,
+    run_base_dir: str,
+    inputs: Inputs = Inputs.empty(),
+    input_path: str | None = None,
+    checkpoints: Checkpoints | None = None,
+    code_bundle: CodeBundle | None = None,
+    image_cache: ImageCache | None = None,
+    interactive_mode: bool = False,
+) -> Tuple[Optional[Outputs], Optional[Error]]:
+    """
+    This method is used to convert the inputs to native types, and run the task. It assumes you are running
+    in a context tree.
+    """
+    ctx = internal_ctx()
+
+    # Load inputs first to get context
+    if input_path:
+        inputs = await load_inputs(input_path, path_rewrite_config=raw_data_path.path_rewrite)
+
+    # Extract context from inputs
+    custom_context = inputs.context if inputs else {}
+
+    tctx = TaskContext(
+        action=action,
+        checkpoints=checkpoints,
+        code_bundle=code_bundle,
+        input_path=input_path,
+        output_path=output_path,
+        run_base_dir=run_base_dir,
+        version=version,
+        raw_data_path=raw_data_path,
+        compiled_image_cache=image_cache,
+        report=flyte.report.Report(name=action.name),
+        mode="remote" if not ctx.data.task_context else ctx.data.task_context.mode,
+        interactive_mode=interactive_mode,
+        custom_context=custom_context,
+    )
+
+    with ctx.replace_task_context(tctx):
+        inputs_kwargs = await convert_inputs_to_native(inputs, task.native_interface)
+        out, err = await run_task(tctx=tctx, controller=controller, task=task, inputs=inputs_kwargs)
+        if err is not None:
+            return None, convert_from_native_to_error(err)
+        if task.report:
+            await flyte.report.flush.aio()
+        return await convert_from_native_to_outputs(out, task.native_interface, task.name), None
+
+
+async def extract_download_run_upload(
+    task: TaskTemplate,
+    *,
+    action: ActionID,
+    controller: Controller,
+    raw_data_path: RawDataPath,
+    output_path: str,
+    run_base_dir: str,
+    version: str,
+    checkpoints: Checkpoints | None = None,
+    code_bundle: CodeBundle | None = None,
+    input_path: str | None = None,
+    image_cache: ImageCache | None = None,
+    interactive_mode: bool = False,
+):
+    """
+    This method is invoked from the CLI (urun) and is used to run a task. This assumes that the context tree
+    has already been created, and the task has been loaded. It also handles the loading of the task.
+    """
+    t = time.time()
+    logger.warning(f"Task {action.name} started at {t}")
+    outputs, err = await convert_and_run(
+        task=task,
+        input_path=input_path,
+        action=action,
+        controller=controller,
+        raw_data_path=raw_data_path,
+        output_path=output_path,
+        run_base_dir=run_base_dir,
+        version=version,
+        checkpoints=checkpoints,
+        code_bundle=code_bundle,
+        image_cache=image_cache,
+        interactive_mode=interactive_mode,
+    )
+    if err is not None:
+        path = await upload_error(err.err, output_path)
+        logger.error(f"Task {task.name} failed with error: {err}. Uploaded error to {path}")
+        return
+    if outputs is None:
+        logger.info(f"Task {task.name} completed successfully, no outputs")
+        return
+    await upload_outputs(outputs, output_path) if output_path else None
+    logger.warning(f"Task {task.name} completed successfully, uploaded outputs to {output_path} in {time.time() - t}s")

flyte/_internal/runtime/trigger_serde.py

@@ -0,0 +1,160 @@
+import asyncio
+from typing import Union
+
+from flyteidl2.core import interface_pb2, literals_pb2
+from flyteidl2.task import common_pb2, run_pb2, task_definition_pb2
+from google.protobuf import timestamp_pb2, wrappers_pb2
+
+import flyte.types
+from flyte import Cron, FixedRate, Trigger, TriggerTime
+
+
+def _to_schedule(m: Union[Cron, FixedRate], kickoff_arg_name: str | None = None) -> common_pb2.Schedule:
+    if isinstance(m, Cron):
+        return common_pb2.Schedule(
+            cron=common_pb2.Cron(
+                expression=m.expression,
+                timezone=m.timezone,
+            ),
+            kickoff_time_input_arg=kickoff_arg_name,
+        )
+    elif isinstance(m, FixedRate):
+        start_time = None
+        if m.start_time is not None:
+            start_time = timestamp_pb2.Timestamp()
+            start_time.FromDatetime(m.start_time)
+
+        return common_pb2.Schedule(
+            rate=common_pb2.FixedRate(
+                value=m.interval_minutes,
+                unit=common_pb2.FixedRateUnit.FIXED_RATE_UNIT_MINUTE,
+                start_time=start_time,
+            ),
+            kickoff_time_input_arg=kickoff_arg_name,
+        )
+
+
+async def process_default_inputs(
+    default_inputs: dict,
+    task_name: str,
+    task_inputs: interface_pb2.VariableMap,
+    task_default_inputs: list[common_pb2.NamedParameter],
+) -> list[common_pb2.NamedLiteral]:
+    """
+    Process default inputs and convert them to NamedLiteral objects.
+
+    Args:
+        default_inputs: Dictionary of default input values
+        task_name: Name of the task for error messages
+        task_inputs: Task input variable map
+        task_default_inputs: List of default parameters from task
+
+    Returns:
+        List of NamedLiteral objects
+    """
+    keys = []
+    literal_coros = []
+    for k, v in default_inputs.items():
+        if k not in task_inputs.variables:
+            raise ValueError(
+                f"Trigger default input '{k}' must be an input to the task, but not found in task {task_name}. "
+                f"Available inputs: {list(task_inputs.variables.keys())}"
+            )
+        else:
+            literal_coros.append(flyte.types.TypeEngine.to_literal(v, type(v), task_inputs.variables[k].type))
+            keys.append(k)
+
+    final_literals: list[literals_pb2.Literal] = await asyncio.gather(*literal_coros)
+
+    for p in task_default_inputs or []:
+        if p.name not in keys:
+            keys.append(p.name)
+            final_literals.append(p.parameter.default)
+
+    literals: list[common_pb2.NamedLiteral] = []
+    for k, lit in zip(keys, final_literals):
+        literals.append(
+            common_pb2.NamedLiteral(
+                name=k,
+                value=lit,
+            )
+        )
+
+    return literals
+
+
+async def to_task_trigger(
+    t: Trigger,
+    task_name: str,
+    task_inputs: interface_pb2.VariableMap,
+    task_default_inputs: list[common_pb2.NamedParameter],
+) -> task_definition_pb2.TaskTrigger:
+    """
+    Converts a Trigger object to a TaskTrigger protobuf object.
+    Args:
+        t:
+        task_name:
+        task_inputs:
+        task_default_inputs:
+    Returns:
+
+    """
+    env = None
+    if t.env_vars:
+        env = run_pb2.Envs()
+        for k, v in t.env_vars.items():
+            env.values.append(literals_pb2.KeyValuePair(key=k, value=v))
+
+    labels = run_pb2.Labels(values=t.labels) if t.labels else None
+
+    annotations = run_pb2.Annotations(values=t.annotations) if t.annotations else None
+
+    run_spec = run_pb2.RunSpec(
+        overwrite_cache=t.overwrite_cache,
+        envs=env,
+        interruptible=wrappers_pb2.BoolValue(value=t.interruptible) if t.interruptible is not None else None,
+        cluster=t.queue,
+        labels=labels,
+        annotations=annotations,
+    )
+
+    kickoff_arg_name = None
+    default_inputs = {}
+    if t.inputs:
+        for k, v in t.inputs.items():
+            if v is TriggerTime:
+                if k == "trigger_time" and k not in task_inputs.variables:
+                    # the 'trigger_time' input name that by default Triggers look for so it's always added.
+                    # Remove it here by skipping if it's not actually an input to the task
+                    continue
+                kickoff_arg_name = k
+            else:
+                default_inputs[k] = v
+
+    # assert that default_inputs and the kickoff_arg_name are infact in the task inputs
+    if kickoff_arg_name is not None and kickoff_arg_name not in task_inputs.variables:
+        raise ValueError(
+            f"For a scheduled trigger, the TriggerTime input '{kickoff_arg_name}' "
+            f"must be an input to the task, but not found in task {task_name}. "
+            f"Available inputs: {list(task_inputs.variables.keys())}"
+        )
+
+    literals = await process_default_inputs(default_inputs, task_name, task_inputs, task_default_inputs)
+
+    automation = _to_schedule(
+        t.automation,
+        kickoff_arg_name=kickoff_arg_name,
+    )
+
+    return task_definition_pb2.TaskTrigger(
+        name=t.name,
+        spec=task_definition_pb2.TaskTriggerSpec(
+            active=t.auto_activate,
+            run_spec=run_spec,
+            inputs=common_pb2.Inputs(literals=literals),
+        ),
+        automation_spec=common_pb2.TriggerAutomationSpec(
+            type=common_pb2.TriggerAutomationSpecType.TYPE_SCHEDULE,
+            schedule=automation,
+        ),
+    )

flyte/_internal/runtime/types_serde.py

@@ -0,0 +1,54 @@
+from typing import Dict, Optional, TypeVar
+
+from flyteidl2.core import interface_pb2
+
+from flyte.models import NativeInterface
+from flyte.types._type_engine import TypeEngine
+
+T = TypeVar("T")
+
+
+def transform_variable_map(
+    variable_map: Dict[str, type],
+    descriptions: Optional[Dict[str, str]] = None,
+) -> Dict[str, interface_pb2.Variable]:
+    """
+    Given a map of str (names of inputs for instance) to their Python native types, return a map of the name to a
+    Flyte Variable object with that type.
+    """
+    res = {}
+    descriptions = descriptions or {}
+    if variable_map:
+        for k, v in variable_map.items():
+            res[k] = transform_type(v, descriptions.get(k, k))
+    return res
+
+
+def transform_native_to_typed_interface(
+    interface: Optional[NativeInterface],
+) -> Optional[interface_pb2.TypedInterface]:
+    """
+    Transform the given simple python native interface to FlyteIDL's interface
+    """
+    if interface is None:
+        return None
+    input_descriptions: Dict[str, str] = {}
+    output_descriptions: Dict[str, str] = {}
+    if interface.docstring:
+        # Fill in descriptions from docstring in the future
+        input_descriptions = {}
+        output_descriptions = {}
+
+    inputs_map = transform_variable_map(interface.get_input_types(), input_descriptions)
+    outputs_map = transform_variable_map(interface.outputs, output_descriptions)
+    return interface_pb2.TypedInterface(
+        inputs=interface_pb2.VariableMap(variables=inputs_map), outputs=interface_pb2.VariableMap(variables=outputs_map)
+    )
+
+
+def transform_type(x: type, description: Optional[str] = None) -> interface_pb2.Variable:
+    # add artifact handling eventually
+    return interface_pb2.Variable(
+        type=TypeEngine.to_literal_type(x),
+        description=description,
+    )

flyte/_keyring/file.py

@@ -0,0 +1,115 @@
+from base64 import decodebytes, encodebytes
+from configparser import ConfigParser, NoOptionError, NoSectionError
+from pathlib import Path
+from typing import Optional
+
+from keyring.backend import KeyringBackend
+from keyring.errors import PasswordDeleteError
+
+_FLYTE_KEYRING_PATH: Path = Path.home() / ".flyte" / "keyring.cfg"
+
+
+class SimplePlainTextKeyring(KeyringBackend):
+    """
+    Simple plain text keyring for remote notebook environments.
+
+    This backend is only active when running in IPython/Jupyter notebooks.
+    For local development, the system keyring is used instead.
+    """
+
+    @property
+    def priority(self):
+        """
+        Return priority based on whether we're in a notebook environment.
+        Negative priority means this backend will be skipped by keyring.
+        """
+        from flyte._tools import ipython_check
+
+        if ipython_check():
+            # In IPython/Jupyter - use this backend
+            return 0.5
+        else:
+            # Not in IPython - skip this backend, let system keyring handle it
+            return -1
+
+    def get_password(self, service: str, username: str) -> Optional[str]:
+        """Get password."""
+        if not self.file_path.exists():
+            return None
+
+        config = ConfigParser(interpolation=None)
+        config.read(self.file_path, encoding="utf-8")
+
+        try:
+            password_base64 = config.get(service, username).encode("utf-8")
+            return decodebytes(password_base64).decode("utf-8")
+        except (NoOptionError, NoSectionError):
+            return None
+
+    def delete_password(self, service: str, username: str) -> None:
+        """Delete password."""
+        if not self.file_path.exists():
+            raise PasswordDeleteError("Config file does not exist")
+
+        config = ConfigParser(interpolation=None)
+        config.read(self.file_path, encoding="utf-8")
+
+        try:
+            if not config.remove_option(service, username):
+                raise PasswordDeleteError("Password not found")
+        except NoSectionError:
+            raise PasswordDeleteError("Password not found")
+
+        with self.file_path.open("w", encoding="utf-8") as config_file:
+            config.write(config_file)
+
+    def set_password(self, service: str, username: str, password: str) -> None:
+        """Set password."""
+        if not username:
+            raise ValueError("Username must be provided")
+
+        file_path = self._ensure_file_path()
+        value = encodebytes(password.encode("utf-8")).decode("utf-8")
+
+        config = ConfigParser(interpolation=None)
+        config.read(file_path, encoding="utf-8")
+
+        if not config.has_section(service):
+            config.add_section(service)
+
+        config.set(service, username, value)
+
+        with file_path.open("w", encoding="utf-8") as config_file:
+            config.write(config_file)
+
+    def _ensure_file_path(self):
+        self.file_path.parent.mkdir(exist_ok=True, parents=True)
+        if not self.file_path.is_file():
+            self.file_path.touch(0o600)
+        return self.file_path
+
+    @property
+    def file_path(self) -> Path:
+        from flyte._initialize import get_init_config, is_initialized
+        from flyte._logging import logger
+
+        # Only try to use source_config_path if flyte.init() has been called
+        if is_initialized():
+            try:
+                config = get_init_config()
+                config_path = config.source_config_path
+                if config_path and str(config_path.parent.name) == ".flyte":
+                    # if the config is in a .flyte directory, use that as the path
+                    return config_path.parent / "keyring.cfg"
+            except Exception as e:
+                # If anything fails, fall back to default path
+                logger.debug(f"Skipping config-based keyring path due to error: {e}")
+        else:
+            # flyte.init() hasn't been called, use default path
+            logger.debug("flyte.init() not called, using default keyring path")
+
+        # Default path
+        return _FLYTE_KEYRING_PATH
+
+    def __repr__(self):
+        return f"<{self.__class__.__name__}> at {self.file_path}>"