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

flyte/_internal/imagebuild/remote_builder.py

@@ -0,0 +1,409 @@
+import gzip
+import os
+import shutil
+import tarfile
+import tempfile
+import typing
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import TYPE_CHECKING, Optional, Tuple, cast
+from uuid import uuid4
+
+import aiofiles
+from flyteidl2.common import phase_pb2
+
+import flyte
+import flyte.errors
+from flyte import Image, remote
+from flyte._code_bundle._utils import tar_strip_file_attributes
+from flyte._image import (
+    _BASE_REGISTRY,
+    AptPackages,
+    Architecture,
+    Commands,
+    CopyConfig,
+    DockerIgnore,
+    Env,
+    PipOption,
+    PipPackages,
+    PoetryProject,
+    PythonWheels,
+    Requirements,
+    UVProject,
+    UVScript,
+    WorkDir,
+)
+from flyte._internal.imagebuild.image_builder import ImageBuilder, ImageChecker
+from flyte._internal.imagebuild.utils import copy_files_to_context, get_and_list_dockerignore
+from flyte._internal.runtime.task_serde import get_security_context
+from flyte._logging import logger
+from flyte._secret import Secret
+from flyte.remote import ActionOutputs, Run
+
+if TYPE_CHECKING:
+    from flyteidl2.imagebuilder import definition_pb2 as image_definition_pb2
+
+IMAGE_TASK_NAME = "build-image"
+IMAGE_TASK_PROJECT = "system"
+IMAGE_TASK_DOMAIN = "production"
+
+
+class RemoteImageChecker(ImageChecker):
+    _images_client = None
+
+    @classmethod
+    async def image_exists(
+        cls, repository: str, tag: str, arch: Tuple[Architecture, ...] = ("linux/amd64",)
+    ) -> Optional[str]:
+        try:
+            import flyte.remote as remote
+
+            remote.Task.get(
+                name=IMAGE_TASK_NAME,
+                project=IMAGE_TASK_PROJECT,
+                domain=IMAGE_TASK_DOMAIN,
+                auto_version="latest",
+            )
+        except Exception as e:
+            msg = "remote image builder is not enabled. Please contact Union support to enable it."
+            raise flyte.errors.ImageBuildError(msg) from e
+
+        image_name = f"{repository.split('/')[-1]}:{tag}"
+
+        try:
+            from flyteidl2.common.identifier_pb2 import ProjectIdentifier
+            from flyteidl2.imagebuilder import definition_pb2 as image_definition__pb2
+            from flyteidl2.imagebuilder import payload_pb2 as image_payload__pb2
+            from flyteidl2.imagebuilder import service_pb2_grpc as image_service_pb2_grpc
+
+            from flyte._initialize import _get_init_config
+
+            cfg = _get_init_config()
+            if cfg is None:
+                raise ValueError("Init config should not be None")
+            image_id = image_definition__pb2.ImageIdentifier(name=image_name)
+            req = image_payload__pb2.GetImageRequest(
+                id=image_id,
+                organization=cfg.org,
+                project_id=ProjectIdentifier(organization=cfg.org, domain=cfg.domain, name=cfg.project),
+            )
+            if cls._images_client is None:
+                if cfg.client is None:
+                    raise ValueError("remote client should not be None")
+                cls._images_client = image_service_pb2_grpc.ImageServiceStub(cfg.client._channel)
+            resp = await cls._images_client.GetImage(req)
+            logger.warning(f"[blue]Image {resp.image.fqin} found. Skip building.[/blue]")
+            return resp.image.fqin
+        except Exception:
+            logger.warning(f"[blue]Image {image_name} was not found or has expired.[/blue]", extra={"highlight": False})
+            return None
+
+
+class RemoteImageBuilder(ImageBuilder):
+    def get_checkers(self) -> Optional[typing.List[typing.Type[ImageChecker]]]:
+        """Return the image checker."""
+        return [RemoteImageChecker]
+
+    async def build_image(self, image: Image, dry_run: bool = False) -> str:
+        image_name = f"{image.name}:{image._final_tag}"
+        spec, context = await _validate_configuration(image)
+
+        start = datetime.now(timezone.utc)
+        try:
+            entity = await remote.Task.get(
+                name=IMAGE_TASK_NAME,
+                project=IMAGE_TASK_PROJECT,
+                domain=IMAGE_TASK_DOMAIN,
+                auto_version="latest",
+            ).override.aio(secrets=_get_build_secrets_from_image(image))
+        except flyte.errors.ReferenceTaskError:
+            raise flyte.errors.ImageBuildError(
+                "remote image builder is not enabled. Please contact Union support to enable it."
+            )
+
+        logger.warning("[bold blue]🐳 Submitting a new build...[/bold blue]")
+        if image.registry and image.registry != _BASE_REGISTRY:
+            target_image = f"{image.registry}/{image_name}"
+        else:
+            # Use the default system registry in the backend.
+            target_image = image_name
+
+        from flyte._initialize import get_init_config
+
+        cfg = get_init_config()
+        run = cast(
+            Run,
+            await flyte.with_runcontext(
+                project=cfg.project, domain=cfg.domain, cache_lookup_scope="project-domain"
+            ).run.aio(entity, spec=spec, context=context, target_image=target_image),
+        )
+        logger.warning(f"⏳ Waiting for build to finish at: [bold cyan link={run.url}]{run.url}[/bold cyan link]")
+
+        await run.wait.aio(quiet=True)
+        run_details = await run.details.aio()
+
+        elapsed = str(datetime.now(timezone.utc) - start).split(".")[0]
+
+        if run_details.action_details.raw_phase == phase_pb2.ACTION_PHASE_SUCCEEDED:
+            logger.warning(f"[bold green]✅ Build completed in {elapsed}![/bold green]")
+        else:
+            raise flyte.errors.ImageBuildError(f"❌ Build failed in {elapsed} at {run.url}")
+
+        outputs = await run_details.outputs()
+        return _get_fully_qualified_image_name(outputs)
+
+
+async def _validate_configuration(image: Image) -> Tuple[str, Optional[str]]:
+    """Validate the configuration and prepare the spec and context files."""  # Prepare the spec file
+    tmp_path = Path(tempfile.gettempdir()) / str(uuid4())
+    os.makedirs(tmp_path, exist_ok=True)
+
+    context_path = tmp_path / "build.uc-image-builder"
+    context_path.mkdir(exist_ok=True)
+
+    image_idl = _get_layers_proto(image, context_path)
+
+    spec_path = tmp_path / "spec.pb"
+    with spec_path.open("wb") as f:
+        f.write(image_idl.SerializeToString())
+
+    _, spec_url = await remote.upload_file.aio(spec_path)
+
+    if any(context_path.iterdir()):
+        # If there are files in the context directory, upload it
+        tar_path = tmp_path / "context.tar"
+        with tarfile.open(tar_path, "w", dereference=False) as tar:
+            files: typing.List[str] = os.listdir(context_path)
+            for ws_file in files:
+                tar.add(
+                    os.path.join(context_path, ws_file),
+                    recursive=True,
+                    arcname=ws_file,
+                    filter=tar_strip_file_attributes,
+                )
+        context_dst = Path(f"{tar_path!s}.gz")
+        with gzip.GzipFile(filename=context_dst, mode="wb", mtime=0) as gzipped:
+            async with aiofiles.open(tar_path, "rb") as tar_file:
+                content = await tar_file.read()
+                gzipped.write(content)
+
+        context_size = tar_path.stat().st_size
+        if context_size > 5 * 1024 * 1024:
+            logger.warning(
+                f"[yellow]Context size is {context_size / (1024 * 1024):.2f} MB, which is larger than 5 MB. "
+                "Upload and build speed will be impacted.[/yellow]",
+            )
+        _, context_url = await remote.upload_file.aio(context_dst)
+    else:
+        context_url = ""
+
+    return spec_url, context_url
+
+
+def _get_layers_proto(image: Image, context_path: Path) -> "image_definition_pb2.ImageSpec":
+    from flyteidl2.imagebuilder import definition_pb2 as image_definition_pb2
+
+    if image.dockerfile is not None:
+        raise flyte.errors.ImageBuildError(
+            "Custom Dockerfile is not supported with remote image builder.You can use local image builder instead."
+        )
+
+    layers = []
+    for layer in image._layers:
+        secret_mounts = None
+        pip_options = image_definition_pb2.PipOptions()
+
+        if isinstance(layer, PipOption):
+            pip_options = image_definition_pb2.PipOptions(
+                index_url=layer.index_url,
+                extra_index_urls=layer.extra_index_urls,
+                pre=layer.pre,
+                extra_args=layer.extra_args,
+            )
+
+        if hasattr(layer, "secret_mounts"):
+            sc = get_security_context(layer.secret_mounts)
+            secret_mounts = sc.secrets if sc else None
+
+        if isinstance(layer, AptPackages):
+            apt_layer = image_definition_pb2.Layer(
+                apt_packages=image_definition_pb2.AptPackages(
+                    packages=layer.packages,
+                    secret_mounts=secret_mounts,
+                ),
+            )
+            layers.append(apt_layer)
+        elif isinstance(layer, PythonWheels):
+            dst_path = copy_files_to_context(layer.wheel_dir, context_path)
+            wheel_layer = image_definition_pb2.Layer(
+                python_wheels=image_definition_pb2.PythonWheels(
+                    dir=str(dst_path.relative_to(context_path)),
+                    options=pip_options,
+                    secret_mounts=secret_mounts,
+                )
+            )
+            layers.append(wheel_layer)
+
+        elif isinstance(layer, Requirements):
+            dst_path = copy_files_to_context(layer.file, context_path)
+            requirements_layer = image_definition_pb2.Layer(
+                requirements=image_definition_pb2.Requirements(
+                    file=str(dst_path.relative_to(context_path)),
+                    options=pip_options,
+                    secret_mounts=secret_mounts,
+                )
+            )
+            layers.append(requirements_layer)
+        elif isinstance(layer, PipPackages) or isinstance(layer, UVScript):
+            if isinstance(layer, UVScript):
+                from flyte._utils import parse_uv_script_file
+
+                header = parse_uv_script_file(layer.script)
+                if not header.dependencies:
+                    continue
+                packages: typing.Iterable[str] = header.dependencies
+                if header.pyprojects:
+                    layers.append(
+                        image_definition_pb2.Layer(
+                            apt_packages=image_definition_pb2.AptPackages(
+                                packages=["git"],  # To get the version of the project.
+                            ),
+                        )
+                    )
+                    docker_ignore_patterns = get_and_list_dockerignore(image)
+
+                    for pyproject in header.pyprojects:
+                        pyproject_dst = copy_files_to_context(Path(pyproject), context_path, docker_ignore_patterns)
+                        uv_project_layer = image_definition_pb2.Layer(
+                            uv_project=image_definition_pb2.UVProject(
+                                pyproject=str(pyproject_dst.relative_to(context_path)),
+                                uvlock=str(
+                                    copy_files_to_context(Path(pyproject) / "uv.lock", context_path).relative_to(
+                                        context_path
+                                    )
+                                ),
+                                options=pip_options,
+                                secret_mounts=secret_mounts,
+                            )
+                        )
+                        layers.append(uv_project_layer)
+
+            else:
+                packages = layer.packages or []
+            pip_layer = image_definition_pb2.Layer(
+                pip_packages=image_definition_pb2.PipPackages(
+                    packages=packages,
+                    options=pip_options,
+                    secret_mounts=secret_mounts,
+                )
+            )
+            layers.append(pip_layer)
+        elif isinstance(layer, UVProject):
+            if layer.project_install_mode == "dependencies_only":
+                # Copy pyproject itself
+                pyproject_dst = copy_files_to_context(layer.pyproject, context_path)
+                if pip_options.extra_args:
+                    if "--no-install-project" not in pip_options.extra_args:
+                        pip_options.extra_args += " --no-install-project"
+                else:
+                    pip_options.extra_args = " --no-install-project"
+                if "--no-sources" not in pip_options.extra_args:
+                    pip_options.extra_args += " --no-sources"
+            else:
+                # Copy the entire project
+                docker_ignore_patterns = get_and_list_dockerignore(image)
+                pyproject_dst = copy_files_to_context(layer.pyproject.parent, context_path, docker_ignore_patterns)
+
+            uv_layer = image_definition_pb2.Layer(
+                uv_project=image_definition_pb2.UVProject(
+                    pyproject=str(pyproject_dst.relative_to(context_path)),
+                    uvlock=str(copy_files_to_context(layer.uvlock, context_path).relative_to(context_path)),
+                    options=pip_options,
+                    secret_mounts=secret_mounts,
+                )
+            )
+            layers.append(uv_layer)
+        elif isinstance(layer, PoetryProject):
+            extra_args = layer.extra_args or ""
+            if layer.project_install_mode == "dependencies_only":
+                # Copy pyproject itself
+                if "--no-root" not in extra_args:
+                    extra_args += " --no-root"
+                pyproject_dst = copy_files_to_context(layer.pyproject, context_path)
+            else:
+                # Copy the entire project
+                pyproject_dst = copy_files_to_context(layer.pyproject.parent, context_path)
+
+            poetry_layer = image_definition_pb2.Layer(
+                poetry_project=image_definition_pb2.PoetryProject(
+                    pyproject=str(pyproject_dst.relative_to(context_path)),
+                    poetry_lock=str(copy_files_to_context(layer.poetry_lock, context_path).relative_to(context_path)),
+                    extra_args=extra_args,
+                    secret_mounts=secret_mounts,
+                )
+            )
+            layers.append(poetry_layer)
+        elif isinstance(layer, Commands):
+            commands_layer = image_definition_pb2.Layer(
+                commands=image_definition_pb2.Commands(
+                    cmd=list(layer.commands),
+                    secret_mounts=secret_mounts,
+                )
+            )
+            layers.append(commands_layer)
+        elif isinstance(layer, DockerIgnore):
+            shutil.copy(layer.path, context_path)
+        elif isinstance(layer, CopyConfig):
+            dst_path = copy_files_to_context(layer.src, context_path)
+
+            copy_layer = image_definition_pb2.Layer(
+                copy_config=image_definition_pb2.CopyConfig(
+                    src=str(dst_path.relative_to(context_path)),
+                    dst=str(layer.dst),
+                )
+            )
+            layers.append(copy_layer)
+        elif isinstance(layer, Env):
+            env_layer = image_definition_pb2.Layer(
+                env=image_definition_pb2.Env(
+                    env_variables=dict(layer.env_vars),
+                )
+            )
+            layers.append(env_layer)
+        elif isinstance(layer, WorkDir):
+            workdir_layer = image_definition_pb2.Layer(
+                workdir=image_definition_pb2.WorkDir(workdir=layer.workdir),
+            )
+            layers.append(workdir_layer)
+
+    return image_definition_pb2.ImageSpec(
+        base_image=image.base_image,
+        python_version=f"{image.python_version[0]}.{image.python_version[1]}",
+        layers=layers,
+    )
+
+
+def _get_fully_qualified_image_name(outputs: ActionOutputs) -> str:
+    return outputs.pb2.literals[0].value.scalar.primitive.string_value
+
+
+def _get_build_secrets_from_image(image: Image) -> Optional[typing.List[Secret]]:
+    secrets = []
+    DEFAULT_SECRET_DIR = Path("/etc/flyte/secrets")
+    for layer in image._layers:
+        if isinstance(layer, (PipOption, Commands, AptPackages)) and layer.secret_mounts is not None:
+            for secret_mount in layer.secret_mounts:
+                # Mount all the image secrets to a default directory that will be passed to the BuildKit server.
+                if isinstance(secret_mount, Secret):
+                    secrets.append(Secret(key=secret_mount.key, group=secret_mount.group, mount=DEFAULT_SECRET_DIR))
+                elif isinstance(secret_mount, str):
+                    secrets.append(Secret(key=secret_mount, mount=DEFAULT_SECRET_DIR))
+                else:
+                    raise ValueError(f"Unsupported secret_mount type: {type(secret_mount)}")
+
+    image_registry_secret = image._image_registry_secret
+    if image_registry_secret:
+        secrets.append(
+            Secret(key=image_registry_secret.key, group=image_registry_secret.group, mount=DEFAULT_SECRET_DIR)
+        )
+    return secrets

flyte/_internal/imagebuild/utils.py

@@ -0,0 +1,79 @@
+import shutil
+from pathlib import Path, PurePath
+from typing import List, Optional
+
+from flyte._image import DockerIgnore, Image
+from flyte._logging import logger
+
+
+def copy_files_to_context(src: Path, context_path: Path, ignore_patterns: list[str] = []) -> Path:
+    """
+    This helper function ensures that absolute paths that users specify are converted correctly to a path in the
+    context directory. Doing this prevents collisions while ensuring files are available in the context.
+
+    For example, if a user has
+        img.with_requirements(Path("/Users/username/requirements.txt"))
+           .with_requirements(Path("requirements.txt"))
+           .with_requirements(Path("../requirements.txt"))
+
+    copying with this function ensures that the Docker context folder has all three files.
+
+    :param src: The source path to copy
+    :param context_path: The context path where the files should be copied to
+    """
+    if src.is_absolute() or ".." in str(src):
+        rel_path = PurePath(*src.parts[1:])
+        dst_path = context_path / "_flyte_abs_context" / rel_path
+    else:
+        dst_path = context_path / src
+    dst_path.parent.mkdir(parents=True, exist_ok=True)
+    if src.is_dir():
+        default_ignore_patterns = [".idea", ".venv"]
+        ignore_patterns = list(set(ignore_patterns + default_ignore_patterns))
+        shutil.copytree(src, dst_path, dirs_exist_ok=True, ignore=shutil.ignore_patterns(*ignore_patterns))
+    else:
+        shutil.copy(src, dst_path)
+    return dst_path
+
+
+def get_and_list_dockerignore(image: Image) -> List[str]:
+    """
+    Get and parse dockerignore patterns from .dockerignore file.
+
+    This function first looks for a DockerIgnore layer in the image's layers. If found, it uses
+    the path specified in that layer. If no DockerIgnore layer is found, it falls back to looking
+    for a .dockerignore file in the root_path directory.
+
+    :param image: The Image object
+    """
+    from flyte._initialize import _get_init_config
+
+    # Look for DockerIgnore layer in the image layers
+    dockerignore_path: Optional[Path] = None
+    patterns: List[str] = []
+
+    for layer in image._layers:
+        if isinstance(layer, DockerIgnore) and layer.path.strip():
+            dockerignore_path = Path(layer.path)
+    # If DockerIgnore layer not specified, set dockerignore_path under root_path
+    init_config = _get_init_config()
+    root_path = init_config.root_dir if init_config else None
+    if not dockerignore_path and root_path:
+        dockerignore_path = Path(root_path) / ".dockerignore"
+    # Return empty list if no .dockerignore file found
+    if not dockerignore_path or not dockerignore_path.exists() or not dockerignore_path.is_file():
+        logger.info(f".dockerignore file not found at path: {dockerignore_path}")
+        return patterns
+
+    try:
+        with open(dockerignore_path, "r", encoding="utf-8") as f:
+            for line in f:
+                stripped_line = line.strip()
+                # Skip empty lines, whitespace-only lines, and comments
+                if not stripped_line or stripped_line.startswith("#"):
+                    continue
+                patterns.append(stripped_line)
+    except Exception as e:
+        logger.error(f"Failed to read .dockerignore file at {dockerignore_path}: {e}")
+        return []
+    return patterns

flyte/_internal/resolvers/_app_env_module.py

@@ -0,0 +1,92 @@
+import pathlib
+import sys
+
+from flyte._logging import logger
+from flyte._module import extract_obj_module
+from flyte.app._app_environment import AppEnvironment
+
+
+def extract_app_env_module(app_env: AppEnvironment, /, source_dir: pathlib.Path) -> tuple[str, str]:
+    """
+    Extract the module name and variable name for a AppEnvironment instance.
+
+    Args:
+        app_env: The AppEnvironment instance to locate.
+        caller_frame: Frame information from where AppEnvironment was instantiated.
+                     If None, falls back to extract_obj_module (which may not work correctly).
+        serialization_context: Context containing the root directory for calculating
+                              relative module paths.
+
+    Returns:
+        A tuple of (module_name, variable_name) where:
+        - module_name: Dotted module path (e.g., "examples.apps.single_script_fastapi")
+        - variable_name: The name of the variable holding the AppEnvironment (e.g., "env")
+
+    Raises:
+        RuntimeError: If the module cannot be loaded or the app variable cannot be found.
+
+    Example:
+        >>> frame = inspect.getframeinfo(inspect.currentframe().f_back)
+        >>> module_name, var_name = _extract_app_env_module_and_var(
+        ...     app_env, frame, serialization_context
+        ... )
+        >>> # Returns: ("examples.apps.my_app", "env")
+        >>> # Can be used as: fserve examples.apps.my_app:env
+    """
+    if app_env._caller_frame is None:
+        raise RuntimeError("Caller frame cannot be None")
+
+    # Get the file path where the app was defined
+    file_path = pathlib.Path(app_env._caller_frame.filename)
+
+    # Calculate module name relative to source_dir
+    try:
+        relative_path = file_path.relative_to(source_dir or pathlib.Path("."))
+        logger.info(f"Relative path: {relative_path}, {source_dir} {pathlib.Path('.')}")
+        module_name = pathlib.Path(relative_path).with_suffix("").as_posix().replace("/", ".")
+    except ValueError:
+        # File is not relative to source_dir, use the stem
+        module_name = file_path.stem
+
+    # Instead of reloading the module, inspect the caller frame's local variables
+    # The app variable should be in the frame's globals
+    caller_globals = None
+
+    # Try to get globals from the main module if it matches our file
+    if hasattr(sys.modules.get("__main__"), "__file__"):
+        main_file = pathlib.Path(sys.modules["__main__"].__file__ or ".").resolve()
+        if main_file == file_path.resolve():
+            caller_globals = sys.modules["__main__"].__dict__
+
+    if caller_globals is None:
+        # Load the module to inspect it by importing it by name
+        # Note: we can't use extract_obj_module here because it uses inspect.getmodule()
+        # which returns the module where the CLASS is defined, not where the INSTANCE is created
+        import importlib
+
+        try:
+            entity_module = importlib.import_module(module_name)
+            caller_globals = entity_module.__dict__
+        except (ModuleNotFoundError, ImportError):
+            # Fallback for test scenarios where module might be <string> or not importable
+            # In this case, use extract_obj_module as a last resort
+            _, entity_module = extract_obj_module(app_env, source_dir)
+            caller_globals = entity_module.__dict__
+
+    # Extract variable name from module - look for AppEnvironment instances
+    app_var_name = None
+    for var_name, obj in caller_globals.items():
+        if isinstance(obj, AppEnvironment):
+            # Found a AppEnvironment - this is likely the one we want
+            # Store the first one we find
+            if app_var_name is None:
+                app_var_name = var_name
+            # If the objects match by identity, use this one
+            if obj is app_env:
+                app_var_name = var_name
+                break
+
+    if app_var_name is None:
+        raise RuntimeError("Could not find variable name for FastAPI app in module")
+
+    return app_var_name, module_name

flyte/_internal/resolvers/_task_module.py

@@ -1,13 +1,11 @@
-import inspect
-import os
 import pathlib
-import sys
 from typing import Tuple
 
+from flyte._module import extract_obj_module
 from flyte._task import AsyncFunctionTaskTemplate, TaskTemplate
 
 
-def extract_task_module(task: TaskTemplate, /, source_dir: pathlib.Path | None = None) -> Tuple[str, str]:
+def extract_task_module(task: TaskTemplate, /, source_dir: pathlib.Path) -> Tuple[str, str]:
     """
     Extract the task module from the task template.
 
@@ -15,40 +13,9 @@ def extract_task_module(task: TaskTemplate, /, source_dir: pathlib.Path | None =
     :param source_dir: The source directory to use for relative paths.
     :return: A tuple containing the entity name, module
     """
-    entity_name = task.name
     if isinstance(task, AsyncFunctionTaskTemplate):
-        entity_module = inspect.getmodule(task.func)
-        if entity_module is None:
-            raise ValueError(f"Task {entity_name} has no module.")
-
-        fp = entity_module.__file__
-        if fp is None:
-            raise ValueError(f"Task {entity_name} has no module.")
-
-        file_path = pathlib.Path(fp)
-        # Get the relative path to the current directory
-        # Will raise ValueError if the file is not in the source directory
-        relative_path = file_path.relative_to(str(source_dir))
-
-        if relative_path == pathlib.Path("."):
-            entity_module_name = entity_module.__name__
-        else:
-            # Replace file separators with dots and remove the '.py' extension
-            dotted_path = os.path.splitext(str(relative_path))[0].replace(os.sep, ".")
-            entity_module_name = dotted_path
-
         entity_name = task.func.__name__
+        entity_module_name, _ = extract_obj_module(task.func, source_dir)
+        return entity_name, entity_module_name
     else:
-        raise NotImplementedError(f"Task module {entity_name} not implemented.")
-
-    if entity_module_name == "__main__":
-        """
-        This case is for the case in which the task is run from the main module.
-        """
-        fp = sys.modules["__main__"].__file__
-        if fp is None:
-            raise ValueError(f"Task {entity_name} has no module.")
-        main_path = pathlib.Path(fp)
-        entity_module_name = main_path.stem
-
-    return entity_name, entity_module_name
+        raise NotImplementedError(f"Task module {task.name} not implemented.")

flyte/_internal/resolvers/app_env.py

@@ -0,0 +1,26 @@
+import importlib
+from pathlib import Path
+
+from flyte._internal.resolvers._app_env_module import extract_app_env_module
+from flyte._internal.resolvers.common import Resolver
+from flyte.app._app_environment import AppEnvironment
+
+
+class AppEnvResolver(Resolver):
+    """
+    Please see the notes in the TaskResolverMixin as it describes this default behavior.
+    """
+
+    @property
+    def import_path(self) -> str:
+        return "flyte._internal.resolvers.app_env.AppEnvResolver"
+
+    def load_app_env(self, loader_args: str) -> AppEnvironment:
+        module_name, app_var_name = loader_args.split(":")
+        app_env_module = importlib.import_module(name=module_name)  # type: ignore
+        app_env_def = getattr(app_env_module, app_var_name)
+        return app_env_def
+
+    def loader_args(self, app_env: AppEnvironment, root_dir: Path) -> str:  # type:ignore
+        app_var_name, module_name = extract_app_env_module(app_env, root_dir)
+        return f"{module_name}:{app_var_name}"

flyte/_internal/resolvers/common.py

@@ -3,6 +3,7 @@ from pathlib import Path
 from typing import List, Optional
 
 from flyte._task import TaskTemplate
+from flyte.app._app_environment import AppEnvironment
 
 
 class Resolver(Protocol):
@@ -23,7 +24,13 @@ class Resolver(Protocol):
         """
         raise NotImplementedError
 
-    def loader_args(self, t: TaskTemplate, root_dir: Optional[Path]) -> List[str]:
+    def load_app_env(self, loader_args: str) -> AppEnvironment:
+        """
+        Given the set of identifier keys, should return one AppEnvironment or raise an error if not found
+        """
+        raise NotImplementedError
+
+    def loader_args(self, t: TaskTemplate, root_dir: Optional[Path]) -> List[str] | str:
         """
         Return a list of strings that can help identify the parameter TaskTemplate. Each string should not have
         spaces or special characters. This is used to identify the task in the resolver.