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

flyte/__init__.py

@@ -9,15 +9,16 @@ 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 ._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
@@ -60,16 +61,20 @@ def version() -> str:
 
 
 __all__ = [
+    "AMD_GPU",
     "GPU",
+    "HABANA_GAUDI",
     "TPU",
     "Cache",
     "CachePolicy",
     "CacheRequest",
     "Cron",
     "Device",
+    "DeviceClass",
     "Environment",
     "FixedRate",
     "Image",
+    "Neuron",
     "PodTemplate",
     "Resources",
     "RetryStrategy",
@@ -85,12 +90,16 @@ __all__ = [
     "build",
     "build_images",
     "ctx",
+    "current_domain",
+    "custom_context",
     "deploy",
+    "get_custom_context",
     "group",
     "init",
     "init_from_config",
     "map",
     "run",
     "trace",
+    "version",
     "with_runcontext",
 ]

flyte/_cache/local_cache.py

@@ -8,9 +8,10 @@ try:
 except ImportError:
     HAS_AIOSQLITE = False
 
+from flyteidl2.task import common_pb2
+
 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"
@@ -127,7 +128,7 @@ class LocalTaskCache(object):
             row = await cursor.fetchone()
             if row:
                 outputs_bytes = row[0]
-                outputs = run_definition_pb2.Outputs()
+                outputs = common_pb2.Outputs()
                 outputs.ParseFromString(outputs_bytes)
                 return convert.Outputs(proto_outputs=outputs)
         return None
@@ -142,7 +143,7 @@ class LocalTaskCache(object):
         row = cursor.fetchone()
         if row:
             outputs_bytes = row[0]
-            outputs = run_definition_pb2.Outputs()
+            outputs = common_pb2.Outputs()
             outputs.ParseFromString(outputs_bytes)
             return convert.Outputs(proto_outputs=outputs)
         return None

flyte/_code_bundle/_utils.py

@@ -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,9 @@ 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)
+        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
@@ -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,17 +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",
-            "--overwrite",
+        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/_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/_deploy.py

@@ -14,13 +14,14 @@ from flyte.syncify import syncify
 
 from ._environment import Environment
 from ._image import Image
-from ._initialize import ensure_client, get_client, get_common_config, requires_initialization
+from ._initialize import ensure_client, get_client, get_init_config, requires_initialization
 from ._logging import logger
 from ._task import TaskTemplate
 from ._task_environment import TaskEnvironment
 
 if TYPE_CHECKING:
-    from flyte._protos.workflow import task_definition_pb2, trigger_definition_pb2
+    from flyteidl2.task import task_definition_pb2
+    from flyteidl2.trigger import trigger_definition_pb2
 
     from ._code_bundle import CopyFiles
     from ._internal.imagebuild.image_builder import ImageCache
@@ -144,11 +145,11 @@ async def _deploy_task(
     """
     ensure_client()
     import grpc.aio
+    from flyteidl2.task import task_definition_pb2, task_service_pb2
 
     from ._internal.runtime.convert import convert_upload_default_inputs
     from ._internal.runtime.task_serde import translate_task_to_wire
     from ._internal.runtime.trigger_serde import to_task_trigger
-    from ._protos.workflow import task_definition_pb2, task_service_pb2
 
     image_uri = task.image.uri if isinstance(task.image, Image) else task.image
 
@@ -213,20 +214,41 @@ async def _build_image_bg(env_name: str, image: Image) -> Tuple[str, str]:
     return env_name, await build.aio(image)
 
 
-async def _build_images(deployment: DeploymentPlan) -> ImageCache:
+async def _build_images(deployment: DeploymentPlan, image_refs: Dict[str, str] | None = None) -> ImageCache:
     """
     Build the images for the given deployment plan and update the environment with the built image.
     """
     from ._internal.imagebuild.image_builder import ImageCache
 
+    if image_refs is None:
+        image_refs = {}
+
     images = []
     image_identifier_map = {}
     for env_name, env in deployment.envs.items():
         if not isinstance(env.image, str):
+            if env.image._ref_name is not None:
+                if env.image._ref_name in image_refs:
+                    # If the image is set in the config, set it as the base_image
+                    image_uri = image_refs[env.image._ref_name]
+                    env.image = env.image.clone(base_image=image_uri)
+                else:
+                    raise ValueError(
+                        f"Image name '{env.image._ref_name}' not found in config. Available: {list(image_refs.keys())}"
+                    )
+                if not env.image._layers:
+                    # No additional layers, use the base_image directly without building
+                    image_identifier_map[env_name] = image_uri
+                    continue
             logger.debug(f"Building Image for environment {env_name}, image: {env.image}")
             images.append(_build_image_bg(env_name, env.image))
 
         elif env.image == "auto" and "auto" not in image_identifier_map:
+            if "default" in image_refs:
+                # If the default image is set through CLI, use it instead
+                image_uri = image_refs["default"]
+                image_identifier_map[env_name] = image_uri
+                continue
             auto_image = Image.from_debian_base()
             images.append(_build_image_bg(env_name, auto_image))
     final_images = await asyncio.gather(*images)
@@ -312,9 +334,9 @@ def get_deployer(env_type: Type[Environment | TaskEnvironment]) -> Deployer:
 async def apply(deployment_plan: DeploymentPlan, copy_style: CopyFiles, dryrun: bool = False) -> Deployment:
     from ._code_bundle import build_code_bundle
 
-    cfg = get_common_config()
+    cfg = get_init_config()
 
-    image_cache = await _build_images(deployment_plan)
+    image_cache = await _build_images(deployment_plan, cfg.images)
 
     if copy_style == "none" and not deployment_plan.version:
         raise flyte.errors.DeploymentError("Version must be set when copy_style is none")
@@ -422,5 +444,7 @@ async def build_images(envs: Environment) -> ImageCache:
     :param envs: Environment to build images for.
     :return: ImageCache containing the built images.
     """
+    cfg = get_init_config()
+    images = cfg.images if cfg else {}
     deployment = plan_deploy(envs)
-    return await _build_images(deployment[0])
+    return await _build_images(deployment[0], images)

flyte/_image.py

@@ -12,8 +12,6 @@ from typing import TYPE_CHECKING, ClassVar, Dict, List, Literal, Optional, Tuple
 import rich.repr
 from packaging.version import Version
 
-from flyte._utils import update_hasher_for_source
-
 if TYPE_CHECKING:
     from flyte import Secret, SecretRequest
 
@@ -167,21 +165,22 @@ class Requirements(PipPackages):
 class UVProject(PipOption, Layer):
     pyproject: Path
     uvlock: Path
+    project_install_mode: typing.Literal["dependencies_only", "install_project"] = "dependencies_only"
 
     def validate(self):
         if not self.pyproject.exists():
-            raise FileNotFoundError(f"pyproject.toml file {self.pyproject} does not exist")
+            raise FileNotFoundError(f"pyproject.toml file {self.pyproject.resolve()} does not exist")
         if not self.pyproject.is_file():
-            raise ValueError(f"Pyproject file {self.pyproject} is not a file")
+            raise ValueError(f"Pyproject file {self.pyproject.resolve()} is not a file")
         if not self.uvlock.exists():
-            raise ValueError(f"UVLock file {self.uvlock} does not exist")
+            raise ValueError(f"UVLock file {self.uvlock.resolve()} does not exist")
         super().validate()
 
     def update_hash(self, hasher: hashlib._Hash):
-        from ._utils import filehash_update
+        from ._utils import filehash_update, update_hasher_for_source
 
         super().update_hash(hasher)
-        if self.extra_args and "--no-install-project" in self.extra_args:
+        if self.project_install_mode == "dependencies_only":
             filehash_update(self.uvlock, hasher)
             filehash_update(self.pyproject, hasher)
         else:
@@ -210,7 +209,7 @@ class PoetryProject(Layer):
         super().validate()
 
     def update_hash(self, hasher: hashlib._Hash):
-        from ._utils import filehash_update
+        from ._utils import filehash_update, update_hasher_for_source
 
         hash_input = ""
         if self.extra_args:
@@ -401,6 +400,8 @@ class Image:
     name: Optional[str] = field(default=None)
     platform: Tuple[Architecture, ...] = field(default=("linux/amd64",))
     python_version: Tuple[int, int] = field(default_factory=_detect_python_version)
+    # Refer to the image_refs (name:image-uri) set in CLI or config
+    _ref_name: Optional[str] = field(default=None)
 
     # Layers to be added to the image. In init, because frozen, but users shouldn't access, so underscore.
     _layers: Tuple[Layer, ...] = field(default_factory=tuple)
@@ -418,6 +419,9 @@ class Image:
     # class-level token not included in __init__
     _token: ClassVar[object] = object()
 
+    # Underscore cuz we may rename in the future, don't expose for now,
+    _image_registry_secret: Optional[Secret] = None
+
     # check for the guard that we put in place
     def __post_init__(self):
         if object.__getattribute__(self, "__dict__").pop("_guard", None) is not Image._token:
@@ -506,6 +510,7 @@ class Image:
         flyte_version: Optional[str] = None,
         install_flyte: bool = True,
         registry: Optional[str] = None,
+        registry_secret: Optional[str | Secret] = None,
         name: Optional[str] = None,
         platform: Optional[Tuple[Architecture, ...]] = None,
     ) -> Image:
@@ -517,6 +522,7 @@ class Image:
         :param flyte_version: Union version to use
         :param install_flyte: If True, will install the flyte library in the image
         :param registry: Registry to use for the image
+        :param registry_secret: Secret to use to pull/push the private image.
         :param name: Name of the image if you want to override the default name
         :param platform: Platform to use for the image, default is linux/amd64, use tuple for multiple values
             Example: ("linux/amd64", "linux/arm64")
@@ -534,7 +540,7 @@ class Image:
         )
 
         if registry or name:
-            return base_image.clone(registry=registry, name=name)
+            return base_image.clone(registry=registry, name=name, registry_secret=registry_secret)
 
         return base_image
 
@@ -549,6 +555,13 @@ class Image:
         img = cls._new(base_image=image_uri)
         return img
 
+    @classmethod
+    def from_ref_name(cls, name: str) -> Image:
+        # NOTE: set image name as _ref_name to enable adding additional layers.
+        # See: https://github.com/flyteorg/flyte-sdk/blob/14de802701aab7b8615ffb99c650a36305ef01f7/src/flyte/_image.py#L642
+        img = cls._new(name=name, _ref_name=name)
+        return img
+
     @classmethod
     def from_uv_script(
         cls,
@@ -556,6 +569,7 @@ class Image:
         *,
         name: str,
         registry: str | None = None,
+        registry_secret: Optional[str | Secret] = None,
         python_version: Optional[Tuple[int, int]] = None,
         index_url: Optional[str] = None,
         extra_index_urls: Union[str, List[str], Tuple[str, ...], None] = None,
@@ -584,6 +598,7 @@ class Image:
 
         :param name: name of the image
         :param registry: registry to use for the image
+        :param registry_secret: Secret to use to pull/push the private image.
         :param python_version: Python version to use for the image, if not specified, will use the current Python
         version
         :param script: path to the uv script
@@ -609,14 +624,22 @@ class Image:
             secret_mounts=_ensure_tuple(secret_mounts) if secret_mounts else None,
         )
 
-        img = cls.from_debian_base(registry=registry, name=name, python_version=python_version, platform=platform)
+        img = cls.from_debian_base(
+            registry=registry,
+            registry_secret=registry_secret,
+            name=name,
+            python_version=python_version,
+            platform=platform,
+        )
 
         return img.clone(addl_layer=ll)
 
     def clone(
         self,
         registry: Optional[str] = None,
+        registry_secret: Optional[str | Secret] = None,
         name: Optional[str] = None,
+        base_image: Optional[str] = None,
         python_version: Optional[Tuple[int, int]] = None,
         addl_layer: Optional[Layer] = None,
     ) -> Image:
@@ -624,12 +647,14 @@ class Image:
         Use this method to clone the current image and change the registry and name
 
         :param registry: Registry to use for the image
+        :param registry_secret: Secret to use to pull/push the private image.
         :param name: Name of the image
         :param python_version: Python version for the image, if not specified, will use the current Python version
         :param addl_layer: Additional layer to add to the image. This will be added to the end of the layers.
-
         :return:
         """
+        from flyte import Secret
+
         if addl_layer and self.dockerfile:
             # We don't know how to inspect dockerfiles to know what kind it is (OS, python version, uv vs poetry, etc)
             # so there's no guarantee any of the layering logic will work.
@@ -639,19 +664,23 @@ class Image:
             )
         registry = registry if registry else self.registry
         name = name if name else self.name
+        registry_secret = registry_secret if registry_secret else self._image_registry_secret
+        base_image = base_image if base_image else self.base_image
         if addl_layer and (not name):
             raise ValueError(
                 f"Cannot add additional layer {addl_layer} to an image without name. Please first clone()."
             )
         new_layers = (*self._layers, addl_layer) if addl_layer else self._layers
         img = Image._new(
-            base_image=self.base_image,
+            base_image=base_image,
             dockerfile=self.dockerfile,
             registry=registry,
             name=name,
             platform=self.platform,
             python_version=python_version or self.python_version,
             _layers=new_layers,
+            _image_registry_secret=Secret(key=registry_secret) if isinstance(registry_secret, str) else registry_secret,
+            _ref_name=self._ref_name,
         )
 
         return img
@@ -876,17 +905,17 @@ class Image:
         pre: bool = False,
         extra_args: Optional[str] = None,
         secret_mounts: Optional[SecretRequest] = None,
+        project_install_mode: typing.Literal["dependencies_only", "install_project"] = "dependencies_only",
     ) -> Image:
         """
         Use this method to create a new image with the specified uv.lock file layered on top of the current image
         Must have a corresponding pyproject.toml file in the same directory
         Cannot be used in conjunction with conda
 
-        By default, this method copies the entire project into the image,
-         including files such as pyproject.toml, uv.lock, and the src/ directory.
+        By default, this method copies the pyproject.toml and uv.lock files into the image.
 
-        If you prefer not to install the current project, you can pass the extra argument --no-install-project.
-         In this case, the image builder will only copy pyproject.toml and uv.lock into the image.
+        If `project_install_mode` is "install_project", it will also copy directory
+         where the pyproject.toml file is located into the image.
 
         :param pyproject_file: path to the pyproject.toml file, needs to have a corresponding uv.lock file
         :param uvlock: path to the uv.lock file, if not specified, will use the default uv.lock file in the same
@@ -896,6 +925,8 @@ class Image:
         :param pre: whether to allow pre-release versions, default is False
         :param extra_args: extra arguments to pass to pip install, default is None
         :param secret_mounts: list of secret mounts to use for the build process.
+        :param project_install_mode: whether to install the project as a package or
+         only dependencies, default is "dependencies_only"
         :return: Image
         """
         if isinstance(pyproject_file, str):
@@ -909,6 +940,7 @@ class Image:
                 pre=pre,
                 extra_args=extra_args,
                 secret_mounts=_ensure_tuple(secret_mounts) if secret_mounts else None,
+                project_install_mode=project_install_mode,
             )
         )
         return new_image