flyte 0.2.0b1__py3-none-any.whl → 2.0.0b46__py3-none-any.whl

flyte/_code_bundle/_packaging.py

@@ -14,34 +14,40 @@ import typing
 from typing import List, Optional, Tuple, Union
 
 import click
-from rich import print as rich_print
 from rich.tree import Tree
 
-from flyte._logging import logger
+from flyte._logging import _get_console, logger
 
 from ._ignore import Ignore, IgnoreGroup
-from ._utils import CopyFiles, _filehash_update, _pathhash_update, ls_files, tar_strip_file_attributes
+from ._utils import (
+    CopyFiles,
+    _filehash_update,
+    _pathhash_update,
+    ls_files,
+    ls_relative_files,
+    tar_strip_file_attributes,
+)
 
 FAST_PREFIX = "fast"
 FAST_FILEENDING = ".tar.gz"
 
 
 def print_ls_tree(source: os.PathLike, ls: typing.List[str]):
-    click.secho("Files to be copied for fast registration...", fg="bright_blue")
+    logger.info("Files to be copied for fast registration...")
 
     tree_root = Tree(
-        f":open_file_folder: [link file://{source}]{source} (detected source root)",
+        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")
@@ -49,7 +55,12 @@ def print_ls_tree(source: os.PathLike, ls: typing.List[str]):
                 else:
                     current = trees[current_path]
         trees[fpp.parent].add(f"{fpp.name}", guide_style="bold bright_blue")
-    rich_print(tree_root)
+
+    console = _get_console()
+    with console.capture() as capture:
+        console.print(tree_root, overflow="ignore", no_wrap=True, crop=False)
+    logger.info(f"Root directory: [link=file://{source}]{source}[/link]")
+    logger.info(capture.get(), extra={"console": console})
 
 
 def _compress_tarball(source: pathlib.Path, output: pathlib.Path) -> None:
@@ -91,10 +102,30 @@ def list_files_to_bundle(
     ignore = IgnoreGroup(source, *ignores)
 
     ls, ls_digest = ls_files(source, copy_style, deref_symlinks, ignore)
-    logger.debug(f"Hash digest: {ls_digest}")
+    logger.debug(f"Hash of files to be included in the code bundle: {ls_digest}")
     return ls, ls_digest
 
 
+def list_relative_files_to_bundle(
+    relative_paths: tuple[str, ...],
+    source: pathlib.Path,
+) -> typing.Tuple[List[str], str]:
+    """
+    List the files in the relative paths.
+
+    :param relative_paths: The list of relative paths to bundle.
+    :param source: The source directory to package.
+    :param ignores: A list of Ignore classes to use for ignoring files
+    :param copy_style: The copy style to use for the tarball
+    :return: A list of all files to be included in the code bundle and a hexdigest of the included files.
+    """
+    _source = source
+
+    all_files, digest = ls_relative_files(list(relative_paths), source)
+    logger.debug(f"Hash of files to be included in the code bundle: {digest}")
+    return all_files, digest
+
+
 def create_bundle(
     source: pathlib.Path, output_dir: pathlib.Path, ls: List[str], ls_digest: str, deref_symlinks: bool = False
 ) -> Tuple[pathlib.Path, float, float]:

flyte/_code_bundle/_utils.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import glob
 import gzip
 import hashlib
 import importlib.util
@@ -14,7 +15,6 @@ import tempfile
 import typing
 from datetime import datetime, timezone
 from functools import lru_cache
-from pathlib import Path
 from types import ModuleType
 from typing import List, Literal, Optional, Tuple, Union
 
@@ -75,9 +75,9 @@ def compress_scripts(source_path: str, destination: str, modules: List[ModuleTyp
 # intended to be passed as a filter to tarfile.add
 # https://docs.python.org/3/library/tarfile.html#tarfile.TarFile.add
 def tar_strip_file_attributes(tar_info: tarfile.TarInfo) -> tarfile.TarInfo:
-    # set time to epoch timestamp 0, aka 00:00:00 UTC on 1 January 1980
+    # set time to epoch timestamp 0, aka 00:00:00 UTC on 1 January 1981
     # note that when extracting this tarfile, this time will be shown as the modified date
-    tar_info.mtime = datetime(1980, 1, 1, tzinfo=timezone.utc).timestamp()
+    tar_info.mtime = datetime(1981, 1, 1, tzinfo=timezone.utc).timestamp()
 
     # user/group info
     tar_info.uid = 0
@@ -135,13 +135,43 @@ def ls_files(
     return all_files, digest
 
 
+def ls_relative_files(relative_paths: list[str], source_path: pathlib.Path) -> tuple[list[str], str]:
+    relative_paths = list(relative_paths)
+    relative_paths.sort()
+    hasher = hashlib.md5()
+
+    all_files: list[str] = []
+    for file in relative_paths:
+        path = source_path / file
+        if path.is_dir():
+            # Filter out directories, only include files
+            all_files.extend([str(p) for p in path.glob("**/*") if p.is_file()])
+        elif path.is_file():
+            all_files.append(str(path))
+        else:
+            glob_files = glob.glob(str(path))
+            if glob_files:
+                # Filter out directories from glob results
+                all_files.extend([str(f) for f in glob_files if pathlib.Path(f).is_file()])
+            else:
+                raise ValueError(f"File {path} is not a valid file, directory, or glob pattern")
+
+    all_files.sort()
+    for p in all_files:
+        _filehash_update(p, hasher)
+        _pathhash_update(p, hasher)
+
+    digest = hasher.hexdigest()
+    return all_files, digest
+
+
 def _filehash_update(path: Union[os.PathLike, str], hasher: hashlib._Hash) -> None:
     blocksize = 65536
     with open(path, "rb") as f:
-        bytes = f.read(blocksize)
-        while bytes:
-            hasher.update(bytes)
-            bytes = f.read(blocksize)
+        chunk = f.read(blocksize)
+        while chunk:
+            hasher.update(chunk)
+            chunk = f.read(blocksize)
 
 
 def _pathhash_update(path: Union[os.PathLike, str], hasher: hashlib._Hash) -> None:
@@ -157,7 +187,7 @@ def list_all_files(source_path: pathlib.Path, deref_symlinks, ignore_group: Opti
 
     # This is needed to prevent infinite recursion when walking with followlinks
     visited_inodes = set()
-    for root, dirnames, files in source_path.walk(top_down=True, follow_symlinks=deref_symlinks):
+    for root, dirnames, files in os.walk(source_path, topdown=True, followlinks=deref_symlinks):
         dirnames[:] = [d for d in dirnames if d not in EXCLUDE_DIRS]
         if deref_symlinks:
             inode = os.stat(root).st_ino
@@ -168,7 +198,7 @@ def list_all_files(source_path: pathlib.Path, deref_symlinks, ignore_group: Opti
         ff = []
         files.sort()
         for fname in files:
-            abspath = (root / fname).absolute()
+            abspath = (pathlib.Path(root) / fname).absolute()
             # Only consider files that exist (e.g. disregard symlinks that point to non-existent files)
             if not os.path.exists(abspath):
                 logger.info(f"Skipping non-existent file {abspath}")
@@ -194,15 +224,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.
@@ -212,12 +242,12 @@ def list_imported_modules_as_files(source_path: str, modules: List[ModuleType])
     import flyte
     from flyte._utils.lazy_module import is_imported
 
-    files = []
-    union_root = os.path.dirname(flyte.__file__)
+    files = set()
+    flyte_root = os.path.dirname(flyte.__file__)
 
     # These directories contain installed packages or modules from the Python standard library.
     # If a module is from these directories, then they are not user files.
-    invalid_directories = [union_root, sys.prefix, sys.base_prefix, site.getusersitepackages(), *site.getsitepackages()]
+    invalid_directories = [flyte_root, sys.prefix, sys.base_prefix, site.getusersitepackages(), *site.getsitepackages()]
 
     for mod in modules:
         # Be careful not to import a module with the .__file__ call if not yet imported.
@@ -241,11 +271,19 @@ def list_imported_modules_as_files(source_path: str, modules: List[ModuleType])
 
         if not _file_is_in_directory(mod_file, source_path):
             # Only upload files where the module file in the source directory
+            # print log line for files that have common ancestor with source_path, but not in it.
+            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
 
-    return files
+        files.add(mod_file)
+
+    return list(files)
 
 
 def add_imported_modules_from_source(source_path: str, destination: str, modules: List[ModuleType]):
@@ -322,18 +360,3 @@ def hash_file(file_path: typing.Union[os.PathLike, str]) -> Tuple[bytes, str, in
             size += len(chunk)
 
     return h.digest(), h.hexdigest(), size
-
-
-def _find_project_root(source_path) -> str:
-    """
-    Find the root of the project.
-    The root of the project is considered to be the first ancestor from source_path that does
-    not contain a __init__.py file.
-
-    N.B.: This assumption only holds for regular packages (as opposed to namespace packages)
-    """
-    # Start from the directory right above source_path
-    path = Path(source_path).parent.resolve()
-    while os.path.exists(os.path.join(path, "__init__.py")):
-        path = path.parent
-    return str(path)

flyte/_code_bundle/bundle.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import asyncio
 import gzip
 import logging
@@ -5,26 +7,54 @@ import os
 import pathlib
 import tempfile
 from pathlib import Path
-from typing import Type
+from typing import TYPE_CHECKING, ClassVar, Type
 
-from flyteidl.core.tasks_pb2 import TaskTemplate
+from async_lru import alru_cache
+from flyteidl2.core.tasks_pb2 import TaskTemplate
 
-import flyte.storage as storage
-from flyte._datastructures import CodeBundle
 from flyte._logging import log, logger
+from flyte._utils import AsyncLRUCache
+from flyte.models import CodeBundle
 
 from ._ignore import GitIgnore, Ignore, StandardIgnore
-from ._packaging import create_bundle, list_files_to_bundle, print_ls_tree
+from ._packaging import create_bundle, list_files_to_bundle, list_relative_files_to_bundle, print_ls_tree
 from ._utils import CopyFiles, hash_file
 
+if TYPE_CHECKING:
+    from flyte.app import AppEnvironment
+
 _pickled_file_extension = ".pkl.gz"
 _tar_file_extension = ".tar.gz"
 
 
+class _PklCache:
+    _pkl_cache: ClassVar[AsyncLRUCache[str, str]] = AsyncLRUCache[str, str](maxsize=100)
+
+    @classmethod
+    async def put(cls, digest: str, upload_to_path: str, from_path: pathlib.Path) -> str:
+        """
+        Get the pickled code bundle from the cache or build it if not present.
+
+        :param digest: The hash digest of the task template.
+        :param upload_to_path: The path to upload the pickled file to.
+        :param from_path: The path to read the pickled file from.
+        :return: CodeBundle object containing the pickled file path and the computed version.
+        """
+        import flyte.storage as storage
+
+        async def put_data() -> str:
+            return await storage.put(str(from_path), to_path=str(upload_to_path))
+
+        return await cls._pkl_cache.get(
+            key=digest,
+            value_func=put_data,
+        )
+
+
 async def build_pkl_bundle(
-    o: TaskTemplate,
+    o: TaskTemplate | AppEnvironment,
     upload_to_controlplane: bool = True,
-    upload_from_dataplane_path: str | None = None,
+    upload_from_dataplane_base_path: str | None = None,
     copy_bundle_to: pathlib.Path | None = None,
 ) -> CodeBundle:
     """
@@ -36,16 +66,14 @@ async def build_pkl_bundle(
 
     :param o: Object to be pickled. This is the task template.
     :param upload_to_controlplane: Whether to upload the pickled file to the control plane or not
-    :param upload_from_dataplane_path: If we are on the dataplane, this is the path where the
+    :param upload_from_dataplane_base_path: If we are on the dataplane, this is the path where the
         pickled file should be uploaded to. upload_to_controlplane has to be False in this case.
     :param copy_bundle_to: If set, the bundle will be copied to this path. This is used for testing purposes.
     :return: CodeBundle object containing the pickled file path and the computed version.
     """
     import cloudpickle
 
-    import flyte.storage as storage
-
-    if upload_to_controlplane and upload_from_dataplane_path:
+    if upload_to_controlplane and upload_from_dataplane_base_path:
         raise ValueError("Cannot upload to control plane and upload from dataplane path at the same time.")
 
     logger.debug("Building pickled code bundle.")
@@ -58,13 +86,20 @@ async def build_pkl_bundle(
             logger.debug("Uploading pickled code bundle to control plane.")
             from flyte.remote import upload_file
 
-            hash_digest, remote_path = await upload_file(dest)
+            hash_digest, remote_path = await upload_file.aio(dest)
             return CodeBundle(pkl=remote_path, computed_version=hash_digest)
 
-        elif upload_from_dataplane_path:
-            logger.debug(f"Uploading pickled code bundle to dataplane path {upload_from_dataplane_path}.")
+        elif upload_from_dataplane_base_path:
+            from flyte._internal.runtime import io
+
             _, str_digest, _ = hash_file(file_path=dest)
-            final_path = await storage.put(str(dest), upload_from_dataplane_path)
+            upload_path = io.pkl_path(upload_from_dataplane_base_path, str_digest)
+            logger.debug(f"Uploading pickled code bundle to dataplane path {upload_path}.")
+            final_path = await _PklCache.put(
+                digest=str_digest,
+                upload_to_path=upload_path,
+                from_path=dest,
+            )
             return CodeBundle(pkl=final_path, computed_version=str_digest)
 
         else:
@@ -74,12 +109,13 @@ 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)
 
 
+@alru_cache
 async def build_code_bundle(
     from_dir: Path,
     *ignore: Type[Ignore],
@@ -90,7 +126,7 @@ async def build_code_bundle(
 ) -> CodeBundle:
     """
     Build the code bundle for the current environment.
-    :param from_dir: The directory to bundle of the code to bundle. This is the root directory for the source.
+    :param from_dir: The directory of the code to bundle. This is the root directory for the source.
     :param extract_dir: The directory to extract the code bundle to, when in the container. It defaults to the current
         working directory.
     :param ignore: The list of ignores to apply. This is a list of Ignore classes.
@@ -111,13 +147,66 @@ async def build_code_bundle(
     if logger.getEffectiveLevel() <= logging.INFO:
         print_ls_tree(from_dir, files)
 
+    logger.debug("Building code bundle.")
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        bundle_path, tar_size, archive_size = create_bundle(
+            from_dir, pathlib.Path(tmp_dir), files, digest, deref_symlinks=True
+        )
+        logger.info(f"Code bundle created at {bundle_path}, size: {tar_size} MB, archive size: {archive_size} MB")
+        if not dryrun:
+            hash_digest, remote_path = await upload_file.aio(bundle_path)
+            logger.debug(f"Code bundle uploaded to {remote_path}")
+        else:
+            if copy_bundle_to:
+                remote_path = str(copy_bundle_to / bundle_path.name)
+            else:
+                import flyte.storage as storage
+
+                base_path = storage.get_random_local_path()
+                base_path.mkdir(parents=True, exist_ok=True)
+                remote_path = str(base_path / bundle_path.name)
+
+            import shutil
+
+            # Copy the bundle to the given path
+            shutil.copy(bundle_path, remote_path)
+            _, hash_digest, _ = hash_file(file_path=bundle_path)
+        return CodeBundle(tgz=remote_path, destination=extract_dir, computed_version=hash_digest, files=files)
+
+
+@alru_cache
+async def build_code_bundle_from_relative_paths(
+    relative_paths: tuple[str, ...],
+    from_dir: Path,
+    extract_dir: str = ".",
+    dryrun: bool = False,
+    copy_bundle_to: pathlib.Path | None = None,
+) -> CodeBundle:
+    """
+    Build a code bundle from a list of relative paths.
+    :param relative_paths: The list of relative paths to bundle.
+    :param from_dir: The directory of the code to bundle. This is the root directory for the source.
+    :param extract_dir: The directory to extract the code bundle to, when in the container. It defaults to the current
+        working directory.
+    :param dryrun: If dryrun is enabled, files will not be uploaded to the control plane.
+    :param copy_bundle_to: If set, the bundle will be copied to this path. This is used for testing purposes.
+    :return: The code bundle, which contains the path where the code was zipped to.
+    """
+    logger.debug("Building code bundle from relative paths.")
+    from flyte.remote import upload_file
+
+    logger.debug("Finding files to bundle")
+    files, digest = list_relative_files_to_bundle(relative_paths, from_dir)
+    if logger.getEffectiveLevel() <= logging.INFO:
+        print_ls_tree(from_dir, files)
+
     logger.debug("Building code bundle.")
     with tempfile.TemporaryDirectory() as tmp_dir:
         bundle_path, tar_size, archive_size = create_bundle(from_dir, pathlib.Path(tmp_dir), files, digest)
         logger.info(f"Code bundle created at {bundle_path}, size: {tar_size} MB, archive size: {archive_size} MB")
         if not dryrun:
-            hash_digest, remote_path = await upload_file(bundle_path)
-            logger.info(f"Code bundle uploaded to {remote_path}")
+            hash_digest, remote_path = await upload_file.aio(bundle_path)
+            logger.debug(f"Code bundle uploaded to {remote_path}")
         else:
             remote_path = "na"
             if copy_bundle_to:
@@ -127,7 +216,7 @@ async def build_code_bundle(
                 shutil.copy(bundle_path, copy_bundle_to)
                 remote_path = str(copy_bundle_to / bundle_path.name)
             _, hash_digest, _ = hash_file(file_path=bundle_path)
-        return CodeBundle(tgz=remote_path, destination=extract_dir, computed_version=hash_digest)
+        return CodeBundle(tgz=remote_path, destination=extract_dir, computed_version=hash_digest, files=files)
 
 
 @log(level=logging.INFO)
@@ -138,29 +227,44 @@ 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)
+    if not dest.exists():
+        dest.mkdir(parents=True, exist_ok=True)
     if not dest.is_dir():
         raise ValueError(f"Destination path should be a directory, found {dest}, {dest.stat()}")
 
     # TODO make storage apis better to accept pathlib.Path
     if bundle.tgz:
         downloaded_bundle = dest / os.path.basename(bundle.tgz)
+        if downloaded_bundle.exists():
+            logger.debug(f"Code bundle {downloaded_bundle} already exists locally, skipping download.")
+            return downloaded_bundle.absolute()
         # Download the tgz file
-        path = await storage.get(bundle.tgz, str(downloaded_bundle.absolute()))
-        downloaded_bundle = pathlib.Path(path)
+        logger.debug(f"Downloading code bundle from {bundle.tgz} to {downloaded_bundle.absolute()}")
+        await storage.get(bundle.tgz, str(downloaded_bundle.absolute()))
         # 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())
@@ -171,8 +275,7 @@ async def download_bundle(bundle: CodeBundle) -> pathlib.Path:
 
         downloaded_bundle = dest / os.path.basename(bundle.pkl)
         # Download the tgz file
-        path = await storage.get(bundle.pkl, str(downloaded_bundle.absolute()))
-        downloaded_bundle = pathlib.Path(path)
+        await storage.get(bundle.pkl, str(downloaded_bundle.absolute()))
         return downloaded_bundle.absolute()
     else:
         raise ValueError("Code bundle should be either tgz or pkl, found neither.")

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

@@ -2,9 +2,10 @@ from __future__ import annotations
 
 import contextvars
 from dataclasses import dataclass, replace
-from typing import TYPE_CHECKING, Awaitable, Callable, Optional, ParamSpec, TypeVar
+from typing import TYPE_CHECKING, Awaitable, Callable, Optional, ParamSpec, Tuple, TypeVar
 
-from flyte._datastructures import GroupData, RawDataPath, TaskContext
+from flyte._logging import logger
+from flyte.models import GroupData, RawDataPath, TaskContext
 
 if TYPE_CHECKING:
     from flyte.report import Report
@@ -28,6 +29,7 @@ class ContextData:
     group_data: Optional[GroupData] = None
     task_context: Optional[TaskContext] = None
     raw_data_path: Optional[RawDataPath] = None
+    metadata: Optional[Tuple[Tuple[str, str], ...]] = None
 
     def replace(self, **kwargs) -> ContextData:
         return replace(self, **kwargs)
@@ -49,6 +51,7 @@ class Context:
             raise ValueError("Cannot create a new context without contextdata.")
         self._data = data
         self._id = id(self)  # Immutable unique identifier
+        self._token = None  # Context variable token to restore the previous context
 
     @property
     def data(self) -> ContextData:
@@ -83,6 +86,12 @@ class Context:
         """
         return Context(self.data.replace(raw_data_path=raw_data_path))
 
+    def new_metadata(self, metadata: Tuple[Tuple[str, str], ...]) -> Context:
+        """
+        Return a copy of the context with the given metadata tuple
+        """
+        return Context(self.data.replace(metadata=metadata))
+
     def get_report(self) -> Optional[Report]:
         """
         Returns a report if within a task context, else a None
@@ -95,7 +104,7 @@ class Context:
     def is_task_context(self) -> bool:
         """
         Returns true if the context is a task context
-        :return:
+        :return: bool
         """
         return self.data.task_context is not None
 
@@ -106,7 +115,11 @@ class Context:
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         """Exit the context, restoring the previous context."""
-        root_context_var.reset(self._token)
+        try:
+            root_context_var.reset(self._token)
+        except Exception as e:
+            logger.warn(f"Failed to reset context: {e}")
+            raise e
 
     async def __aenter__(self):
         """Async version of context entry."""
@@ -129,7 +142,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