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

flyte/__init__.py

@@ -9,15 +9,17 @@ import sys
 from ._build import build
 from ._cache import Cache, CachePolicy, CacheRequest
 from ._context import ctx
+from ._custom_context import custom_context, get_custom_context
 from ._deploy import build_images, deploy
 from ._environment import Environment
 from ._excepthook import custom_excepthook
 from ._group import group
 from ._image import Image
-from ._initialize import init, init_from_config
+from ._initialize import current_domain, init, init_from_config
+from ._logging import logger
 from ._map import map
 from ._pod import PodTemplate
-from ._resources import GPU, TPU, Device, Resources
+from ._resources import AMD_GPU, GPU, HABANA_GAUDI, TPU, Device, DeviceClass, Neuron, Resources
 from ._retry import RetryStrategy
 from ._reusable_environment import ReusePolicy
 from ._run import run, with_runcontext
@@ -25,6 +27,7 @@ from ._secret import Secret, SecretRequest
 from ._task_environment import TaskEnvironment
 from ._timeout import Timeout, TimeoutType
 from ._trace import trace
+from ._trigger import Cron, FixedRate, Trigger, TriggerTime
 from ._version import __version__
 
 sys.excepthook = custom_excepthook
@@ -59,14 +62,20 @@ def version() -> str:
 
 
 __all__ = [
+    "AMD_GPU",
     "GPU",
+    "HABANA_GAUDI",
     "TPU",
     "Cache",
     "CachePolicy",
     "CacheRequest",
+    "Cron",
     "Device",
+    "DeviceClass",
     "Environment",
+    "FixedRate",
     "Image",
+    "Neuron",
     "PodTemplate",
     "Resources",
     "RetryStrategy",
@@ -76,16 +85,23 @@ __all__ = [
     "TaskEnvironment",
     "Timeout",
     "TimeoutType",
+    "Trigger",
+    "TriggerTime",
     "__version__",
     "build",
     "build_images",
     "ctx",
+    "current_domain",
+    "custom_context",
     "deploy",
+    "get_custom_context",
     "group",
     "init",
     "init_from_config",
+    "logger",
     "map",
     "run",
     "trace",
+    "version",
     "with_runcontext",
 ]

flyte/_bin/runtime.py

@@ -12,6 +12,9 @@ from typing import Any, List
 
 import click
 
+from flyte._utils.helpers import str2bool
+from flyte.models import PathRewrite
+
 # Todo: work with pvditt to make these the names
 # ACTION_NAME = "_U_ACTION_NAME"
 # RUN_NAME = "_U_RUN_NAME"
@@ -25,11 +28,12 @@ PROJECT_NAME = "FLYTE_INTERNAL_EXECUTION_PROJECT"
 DOMAIN_NAME = "FLYTE_INTERNAL_EXECUTION_DOMAIN"
 ORG_NAME = "_U_ORG_NAME"
 ENDPOINT_OVERRIDE = "_U_EP_OVERRIDE"
+INSECURE_SKIP_VERIFY_OVERRIDE = "_U_INSECURE_SKIP_VERIFY"
 RUN_OUTPUT_BASE_DIR = "_U_RUN_BASE"
 FLYTE_ENABLE_VSCODE_KEY = "_F_E_VS"
 
-# TODO: Remove this after proper auth is implemented
 _UNION_EAGER_API_KEY_ENV_VAR = "_UNION_EAGER_API_KEY"
+_F_PATH_REWRITE = "_F_PATH_REWRITE"
 
 
 @click.group()
@@ -94,6 +98,7 @@ def main(
     import flyte
     import flyte._utils as utils
     import flyte.errors
+    import flyte.storage as storage
     from flyte._initialize import init
     from flyte._internal.controllers import create_controller
     from flyte._internal.imagebuild.image_builder import ImageCache
@@ -136,19 +141,40 @@ def main(
             controller_kwargs["insecure"] = True
         logger.debug(f"Using controller endpoint: {ep} with kwargs: {controller_kwargs}")
 
-    bundle = CodeBundle(tgz=tgz, pkl=pkl, destination=dest, computed_version=version)
+    # Check for insecure_skip_verify override (e.g. for self-signed certs)
+    insecure_skip_verify_str = os.getenv(INSECURE_SKIP_VERIFY_OVERRIDE, "")
+    if str2bool(insecure_skip_verify_str):
+        controller_kwargs["insecure_skip_verify"] = True
+        logger.info("SSL certificate verification disabled (insecure_skip_verify=True)")
+
+    bundle = None
+    if tgz or pkl:
+        bundle = CodeBundle(tgz=tgz, pkl=pkl, destination=dest, computed_version=version)
     init(org=org, project=project, domain=domain, image_builder="remote", **controller_kwargs)
     # Controller is created with the same kwargs as init, so that it can be used to run tasks
     controller = create_controller(ct="remote", **controller_kwargs)
 
     ic = ImageCache.from_transport(image_cache) if image_cache else None
 
+    path_rewrite_cfg = os.getenv(_F_PATH_REWRITE, None)
+    path_rewrite = None
+    if path_rewrite_cfg:
+        potential_path_rewrite = PathRewrite.from_str(path_rewrite_cfg)
+        if storage.exists_sync(potential_path_rewrite.new_prefix):
+            path_rewrite = potential_path_rewrite
+            logger.info(f"Path rewrite configured for {path_rewrite.new_prefix}")
+        else:
+            logger.error(
+                f"Path rewrite failed for path {potential_path_rewrite.new_prefix}, "
+                f"not found, reverting to original path {potential_path_rewrite.old_prefix}"
+            )
+
     # Create a coroutine to load the task and run it
     task_coroutine = load_and_run_task(
         resolver=resolver,
         resolver_args=resolver_args,
         action=ActionID(name=name, run_name=run_name, project=project, domain=domain, org=org),
-        raw_data_path=RawDataPath(path=raw_data_path),
+        raw_data_path=RawDataPath(path=raw_data_path, path_rewrite=path_rewrite),
         checkpoints=Checkpoints(checkpoint_path, prev_checkpoint),
         code_bundle=bundle,
         input_path=inputs,
@@ -166,8 +192,20 @@ def main(
     async def _run_and_stop():
         loop = asyncio.get_event_loop()
         loop.set_exception_handler(flyte.errors.silence_grpc_polling_error)
-        await utils.run_coros(controller_failure, task_coroutine)
-        await controller.stop()
+        try:
+            await utils.run_coros(controller_failure, task_coroutine)
+            await controller.stop()
+        except flyte.errors.RuntimeSystemError as e:
+            logger.error(f"Runtime system error: {e}")
+            from flyte._internal.runtime.convert import convert_from_native_to_error
+            from flyte._internal.runtime.io import upload_error
+
+            logger.error(f"Flyte runtime failed for action {name} with run name {run_name}, error: {e}")
+            err = convert_from_native_to_error(e)
+            path = await upload_error(err.err, outputs_path)
+            logger.error(f"Run {run_name} Action {name} failed with error: {err}. Uploaded error to {path}")
+            await controller.stop()
+            raise
 
     asyncio.run(_run_and_stop())
     logger.warning(f"Flyte runtime completed for action {name} with run name {run_name}")

flyte/_cache/cache.py

@@ -77,14 +77,16 @@ class Cache:
     def __post_init__(self):
         if self.behavior not in get_args(CacheBehavior):
             raise ValueError(f"Invalid cache behavior: {self.behavior}. Must be one of ['auto', 'override', 'disable']")
-        if self.behavior == "disable":
-            return
 
+        # Still setup _ignore_inputs when cache is disabled to prevent _ignored_inputs attribute not found error
         if isinstance(self.ignored_inputs, str):
             self._ignored_inputs = (self.ignored_inputs,)
         else:
             self._ignored_inputs = self.ignored_inputs
 
+        if self.behavior == "disable":
+            return
+
         # Normalize policies so that self._policies is always a list
         if self.policies is None:
             from flyte._cache.defaults import get_default_policies

flyte/_cache/local_cache.py

@@ -0,0 +1,216 @@
+import sqlite3
+from pathlib import Path
+
+try:
+    import aiosqlite
+
+    HAS_AIOSQLITE = True
+except ImportError:
+    HAS_AIOSQLITE = False
+
+from flyteidl2.task import common_pb2
+
+from flyte._internal.runtime import convert
+from flyte._logging import logger
+from flyte.config import auto
+
+DEFAULT_CACHE_DIR = "~/.flyte"
+CACHE_LOCATION = "local-cache/cache.db"
+
+
+class LocalTaskCache(object):
+    """
+    This class implements a persistent store able to cache the result of local task executions.
+    """
+
+    _conn: "aiosqlite.Connection | None" = None
+    _conn_sync: sqlite3.Connection | None = None
+    _initialized: bool = False
+
+    @staticmethod
+    def _get_cache_path() -> str:
+        """Get the cache database path, creating directory if needed."""
+        config = auto()
+        if config.source:
+            cache_dir = config.source.parent
+        else:
+            cache_dir = Path(DEFAULT_CACHE_DIR).expanduser()
+
+        cache_path = cache_dir / CACHE_LOCATION
+        # Ensure the directory exists
+        cache_path.parent.mkdir(parents=True, exist_ok=True)
+        logger.info(f"Use local cache path: {cache_path}")
+        return str(cache_path)
+
+    @staticmethod
+    async def initialize():
+        """Initialize the cache with database connection."""
+        if not LocalTaskCache._initialized:
+            if HAS_AIOSQLITE:
+                await LocalTaskCache._initialize_async()
+            else:
+                LocalTaskCache._initialize_sync()
+
+    @staticmethod
+    async def _initialize_async():
+        """Initialize async cache connection."""
+        db_path = LocalTaskCache._get_cache_path()
+        conn = await aiosqlite.connect(db_path)
+        await conn.execute("""
+            CREATE TABLE IF NOT EXISTS task_cache (
+                key TEXT PRIMARY KEY,
+                value BLOB
+            )
+        """)
+        await conn.commit()
+        LocalTaskCache._conn = conn
+        LocalTaskCache._initialized = True
+
+    @staticmethod
+    def _initialize_sync():
+        """Initialize sync cache connection."""
+        db_path = LocalTaskCache._get_cache_path()
+        conn = sqlite3.connect(db_path)
+        conn.execute("""
+            CREATE TABLE IF NOT EXISTS task_cache (
+                key TEXT PRIMARY KEY,
+                value BLOB
+            )
+        """)
+        conn.commit()
+        LocalTaskCache._conn_sync = conn
+        LocalTaskCache._initialized = True
+
+    @staticmethod
+    async def clear():
+        """Clear all cache entries."""
+        if not LocalTaskCache._initialized:
+            await LocalTaskCache.initialize()
+
+        if HAS_AIOSQLITE:
+            await LocalTaskCache._clear_async()
+        else:
+            LocalTaskCache._clear_sync()
+
+    @staticmethod
+    async def _clear_async():
+        """Clear all cache entries (async)."""
+        if LocalTaskCache._conn is None:
+            raise RuntimeError("Cache not properly initialized")
+        await LocalTaskCache._conn.execute("DELETE FROM task_cache")
+        await LocalTaskCache._conn.commit()
+
+    @staticmethod
+    def _clear_sync():
+        """Clear all cache entries (sync)."""
+        if LocalTaskCache._conn_sync is None:
+            raise RuntimeError("Cache not properly initialized")
+        LocalTaskCache._conn_sync.execute("DELETE FROM task_cache")
+        LocalTaskCache._conn_sync.commit()
+
+    @staticmethod
+    async def get(cache_key: str) -> convert.Outputs | None:
+        if not LocalTaskCache._initialized:
+            await LocalTaskCache.initialize()
+
+        if HAS_AIOSQLITE:
+            return await LocalTaskCache._get_async(cache_key)
+        else:
+            return LocalTaskCache._get_sync(cache_key)
+
+    @staticmethod
+    async def _get_async(cache_key: str) -> convert.Outputs | None:
+        """Get cache entry (async)."""
+        if LocalTaskCache._conn is None:
+            raise RuntimeError("Cache not properly initialized")
+
+        async with LocalTaskCache._conn.execute("SELECT value FROM task_cache WHERE key = ?", (cache_key,)) as cursor:
+            row = await cursor.fetchone()
+            if row:
+                outputs_bytes = row[0]
+                outputs = common_pb2.Outputs()
+                outputs.ParseFromString(outputs_bytes)
+                return convert.Outputs(proto_outputs=outputs)
+        return None
+
+    @staticmethod
+    def _get_sync(cache_key: str) -> convert.Outputs | None:
+        """Get cache entry (sync)."""
+        if LocalTaskCache._conn_sync is None:
+            raise RuntimeError("Cache not properly initialized")
+
+        cursor = LocalTaskCache._conn_sync.execute("SELECT value FROM task_cache WHERE key = ?", (cache_key,))
+        row = cursor.fetchone()
+        if row:
+            outputs_bytes = row[0]
+            outputs = common_pb2.Outputs()
+            outputs.ParseFromString(outputs_bytes)
+            return convert.Outputs(proto_outputs=outputs)
+        return None
+
+    @staticmethod
+    async def set(
+        cache_key: str,
+        value: convert.Outputs,
+    ) -> None:
+        if not LocalTaskCache._initialized:
+            await LocalTaskCache.initialize()
+
+        if HAS_AIOSQLITE:
+            await LocalTaskCache._set_async(cache_key, value)
+        else:
+            LocalTaskCache._set_sync(cache_key, value)
+
+    @staticmethod
+    async def _set_async(
+        cache_key: str,
+        value: convert.Outputs,
+    ) -> None:
+        """Set cache entry (async)."""
+        if LocalTaskCache._conn is None:
+            raise RuntimeError("Cache not properly initialized")
+
+        output_bytes = value.proto_outputs.SerializeToString()
+        await LocalTaskCache._conn.execute(
+            "INSERT OR REPLACE INTO task_cache (key, value) VALUES (?, ?)", (cache_key, output_bytes)
+        )
+        await LocalTaskCache._conn.commit()
+
+    @staticmethod
+    def _set_sync(
+        cache_key: str,
+        value: convert.Outputs,
+    ) -> None:
+        """Set cache entry (sync)."""
+        if LocalTaskCache._conn_sync is None:
+            raise RuntimeError("Cache not properly initialized")
+
+        output_bytes = value.proto_outputs.SerializeToString()
+        LocalTaskCache._conn_sync.execute(
+            "INSERT OR REPLACE INTO task_cache (key, value) VALUES (?, ?)", (cache_key, output_bytes)
+        )
+        LocalTaskCache._conn_sync.commit()
+
+    @staticmethod
+    async def close():
+        """Close the database connection."""
+        if HAS_AIOSQLITE:
+            await LocalTaskCache._close_async()
+        else:
+            LocalTaskCache._close_sync()
+
+    @staticmethod
+    async def _close_async():
+        """Close async database connection."""
+        if LocalTaskCache._conn:
+            await LocalTaskCache._conn.close()
+            LocalTaskCache._conn = None
+        LocalTaskCache._initialized = False
+
+    @staticmethod
+    def _close_sync():
+        """Close sync database connection."""
+        if LocalTaskCache._conn_sync:
+            LocalTaskCache._conn_sync.close()
+            LocalTaskCache._conn_sync = None
+        LocalTaskCache._initialized = False

flyte/_code_bundle/_ignore.py

@@ -79,7 +79,7 @@ class StandardIgnore(Ignore):
     by fed with custom ignore patterns from cli."""
 
     def __init__(self, root: Path, patterns: Optional[List[str]] = None):
-        super().__init__(root)
+        super().__init__(root.resolve())
         self.patterns = patterns if patterns else STANDARD_IGNORE_PATTERNS
 
     def _is_ignored(self, path: pathlib.Path) -> bool:

flyte/_code_bundle/_packaging.py

@@ -32,15 +32,15 @@ def print_ls_tree(source: os.PathLike, ls: typing.List[str]):
         f"File structure:\n:open_file_folder: {source}",
         guide_style="bold bright_blue",
     )
-    trees = {pathlib.Path(source): tree_root}
-
+    source_path = pathlib.Path(source).resolve()
+    trees = {source_path: tree_root}
     for f in ls:
         fpp = pathlib.Path(f)
         if fpp.parent not in trees:
             # add trees for all intermediate folders
             current = tree_root
-            current_path = pathlib.Path(source)
-            for subdir in fpp.parent.relative_to(source).parts:
+            current_path = source_path  # pathlib.Path(source)
+            for subdir in fpp.parent.relative_to(source_path).parts:
                 current_path = current_path / subdir
                 if current_path not in trees:
                     current = current.add(f"{subdir}", guide_style="bold bright_blue")

flyte/_code_bundle/_utils.py

@@ -193,15 +193,15 @@ def list_all_files(source_path: pathlib.Path, deref_symlinks, ignore_group: Opti
 def _file_is_in_directory(file: str, directory: str) -> bool:
     """Return True if file is in directory and in its children."""
     try:
-        return os.path.commonpath([file, directory]) == directory
-    except ValueError as e:
-        # ValueError is raised by windows if the paths are not from the same drive
-        logger.debug(f"{file} and {directory} are not in the same drive: {e!s}")
+        return pathlib.Path(file).resolve().is_relative_to(pathlib.Path(directory).resolve())
+    except OSError as e:
+        # OSError can be raised if paths cannot be resolved (permissions, broken symlinks, etc.)
+        logger.debug(f"Failed to resolve paths for {file} and {directory}: {e!s}")
         return False
 
 
 def list_imported_modules_as_files(source_path: str, modules: List[ModuleType]) -> List[str]:
-    """Copies modules into destination that are in modules. The module files are copied only if:
+    """Lists the files of modules that have been loaded.  The files are only included if:
 
     1. Not a site-packages. These are installed packages and not user files.
     2. Not in the sys.base_prefix or sys.prefix. These are also installed and not user files.
@@ -211,7 +211,7 @@ def list_imported_modules_as_files(source_path: str, modules: List[ModuleType])
     import flyte
     from flyte._utils.lazy_module import is_imported
 
-    files = []
+    files = set()
     flyte_root = os.path.dirname(flyte.__file__)
 
     # These directories contain installed packages or modules from the Python standard library.
@@ -244,9 +244,15 @@ def list_imported_modules_as_files(source_path: str, modules: List[ModuleType])
             logger.debug(f"{mod_file} is not in {source_path}")
             continue
 
-        files.append(mod_file)
+        if not pathlib.Path(mod_file).is_file():
+            # Some modules have a __file__ attribute that are relative to the base package. Let's skip these,
+            # can add more rigorous logic to really pull out the correct file location if we need to.
+            logger.debug(f"Skipping {mod_file} from {mod.__name__} because it is not a file")
+            continue
+
+        files.add(mod_file)
 
-    return files
+    return list(files)
 
 
 def add_imported_modules_from_source(source_path: str, destination: str, modules: List[ModuleType]):

flyte/_code_bundle/bundle.py

@@ -8,7 +8,7 @@ from pathlib import Path
 from typing import ClassVar, Type
 
 from async_lru import alru_cache
-from flyteidl.core.tasks_pb2 import TaskTemplate
+from flyteidl2.core.tasks_pb2 import TaskTemplate
 
 from flyte._logging import log, logger
 from flyte._utils import AsyncLRUCache
@@ -104,7 +104,7 @@ async def build_pkl_bundle(
                 import shutil
 
                 # Copy the bundle to the given path
-                shutil.copy(dest, copy_bundle_to)
+                shutil.copy(dest, copy_bundle_to, follow_symlinks=True)
                 local_path = copy_bundle_to / dest.name
                 return CodeBundle(pkl=str(local_path), computed_version=str_digest)
             return CodeBundle(pkl=str(dest), computed_version=str_digest)
@@ -169,6 +169,8 @@ async def download_bundle(bundle: CodeBundle) -> pathlib.Path:
 
     :return: The path to the downloaded code bundle.
     """
+    import sys
+
     import flyte.storage as storage
 
     dest = pathlib.Path(bundle.destination)
@@ -185,16 +187,22 @@ async def download_bundle(bundle: CodeBundle) -> pathlib.Path:
         # NOTE the os.path.join(destination, ''). This is to ensure that the given path is in fact a directory and all
         # downloaded data should be copied into this directory. We do this to account for a difference in behavior in
         # fsspec, which requires a trailing slash in case of pre-existing directory.
-        process = await asyncio.create_subprocess_exec(
-            "tar",
+        args = [
             "-xvf",
             str(downloaded_bundle),
             "-C",
             str(dest),
+        ]
+        if sys.platform != "darwin":
+            args.insert(0, "--overwrite")
+
+        process = await asyncio.create_subprocess_exec(
+            "tar",
+            *args,
             stdout=asyncio.subprocess.PIPE,
             stderr=asyncio.subprocess.PIPE,
         )
-        stdout, stderr = await process.communicate()
+        _stdout, stderr = await process.communicate()
 
         if process.returncode != 0:
             raise RuntimeError(stderr.decode())

flyte/_constants.py

@@ -0,0 +1 @@
+FLYTE_SYS_PATH = "_F_SYS_PATH"  # The paths that will be appended to sys.path at runtime

flyte/_context.py

@@ -135,7 +135,10 @@ root_context_var = contextvars.ContextVar("root", default=Context(data=ContextDa
 
 
 def ctx() -> Optional[TaskContext]:
-    """Retrieve the current task context from the context variable."""
+    """
+    Returns flyte.models.TaskContext if within a task context, else None
+    Note: Only use this in task code and not module level.
+    """
     return internal_ctx().data.task_context
 
 

flyte/_custom_context.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+
+from flyte._context import ctx
+
+from ._context import internal_ctx
+
+
+def get_custom_context() -> dict[str, str]:
+    """
+    Get the current input context. This can be used within a task to retrieve
+    context metadata that was passed to the action.
+
+    Context will automatically propagate to sub-actions.
+
+    Example:
+    ```python
+    import flyte
+
+    env = flyte.TaskEnvironment(name="...")
+
+    @env.task
+    def t1():
+        # context can be retrieved with `get_custom_context`
+        ctx = flyte.get_custom_context()
+        print(ctx)  # {'project': '...', 'entity': '...'}
+    ```
+
+    :return: Dictionary of context key-value pairs
+    """
+    tctx = ctx()
+    if tctx is None or tctx.custom_context is None:
+        return {}
+    return tctx.custom_context
+
+
+@contextmanager
+def custom_context(**context: str):
+    """
+    Synchronous context manager to set input context for tasks spawned within this block.
+
+    Example:
+    ```python
+    import flyte
+
+    env = flyte.TaskEnvironment(name="...")
+
+    @env.task
+    def t1():
+        ctx = flyte.get_custom_context()
+        print(ctx)
+
+    @env.task
+    def main():
+        # context can be passed via a context manager
+        with flyte.custom_context(project="my-project"):
+            t1()  # will have {'project': 'my-project'} as context
+    ```
+
+    :param context: Key-value pairs to set as input context
+    """
+    ctx = internal_ctx()
+    if ctx.data.task_context is None:
+        yield
+        return
+
+    tctx = ctx.data.task_context
+    new_tctx = tctx.replace(custom_context={**tctx.custom_context, **context})
+
+    with ctx.replace_task_context(new_tctx):
+        yield
+        # Exit the context and restore the previous context

flyte/_debug/constants.py

@@ -12,7 +12,6 @@ DEFAULT_CODE_SERVER_REMOTE_PATHS = {
 }
 DEFAULT_CODE_SERVER_EXTENSIONS = [
     "https://raw.githubusercontent.com/flyteorg/flytetools/master/flytekitplugins/flyin/ms-python.python-2023.20.0.vsix",
-    "https://raw.githubusercontent.com/flyteorg/flytetools/master/flytekitplugins/flyin/ms-toolsai.jupyter-2023.9.100.vsix",
 ]
 
 # Duration to pause the checking of the heartbeat file until the next one

flyte/_debug/vscode.py

@@ -256,7 +256,12 @@ def prepare_launch_json(ctx: click.Context, pid: int):
 
 
 async def _start_vscode_server(ctx: click.Context):
-    await asyncio.gather(download_tgz(ctx.params["dest"], ctx.params["version"], ctx.params["tgz"]), download_vscode())
+    if ctx.params["tgz"] is None:
+        await download_vscode()
+    else:
+        await asyncio.gather(
+            download_tgz(ctx.params["dest"], ctx.params["version"], ctx.params["tgz"]), download_vscode()
+        )
     child_process = multiprocessing.Process(
         target=lambda cmd: asyncio.run(asyncio.run(execute_command(cmd))),
         kwargs={"cmd": f"code-server --bind-addr 0.0.0.0:6060 --disable-workspace-trust --auth none {os.getcwd()}"},