flyte 2.0.0b21__py3-none-any.whl → 2.0.0b23__py3-none-any.whl

flyte/__init__.py

@@ -25,6 +25,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
@@ -64,8 +65,10 @@ __all__ = [
     "Cache",
     "CachePolicy",
     "CacheRequest",
+    "Cron",
     "Device",
     "Environment",
+    "FixedRate",
     "Image",
     "PodTemplate",
     "Resources",
@@ -76,6 +79,8 @@ __all__ = [
     "TaskEnvironment",
     "Timeout",
     "TimeoutType",
+    "Trigger",
+    "TriggerTime",
     "__version__",
     "build",
     "build_images",

flyte/_bin/runtime.py

@@ -12,6 +12,8 @@ from typing import Any, List
 
 import click
 
+from flyte.models import PathRewrite
+
 # Todo: work with pvditt to make these the names
 # ACTION_NAME = "_U_ACTION_NAME"
 # RUN_NAME = "_U_RUN_NAME"
@@ -28,8 +30,8 @@ ENDPOINT_OVERRIDE = "_U_EP_OVERRIDE"
 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 +96,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 +139,34 @@ 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)
-    init(org=org, project=project, domain=domain, **controller_kwargs)
+    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 +184,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,215 @@
+import sqlite3
+from pathlib import Path
+
+try:
+    import aiosqlite
+
+    HAS_AIOSQLITE = True
+except ImportError:
+    HAS_AIOSQLITE = False
+
+from flyte._internal.runtime import convert
+from flyte._logging import logger
+from flyte._protos.workflow import run_definition_pb2
+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 = run_definition_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 = run_definition_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/bundle.py

@@ -187,6 +187,7 @@ async def download_bundle(bundle: CodeBundle) -> pathlib.Path:
         # fsspec, which requires a trailing slash in case of pre-existing directory.
         process = await asyncio.create_subprocess_exec(
             "tar",
+            "--overwrite",
             "-xvf",
             str(downloaded_bundle),
             "-C",

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()}"},