flyte 2.0.0b9__py3-none-any.whl → 2.0.0b14__py3-none-any.whl

flyte/_image.py

@@ -252,9 +252,15 @@ class AptPackages(Layer):
 @dataclass(frozen=True, repr=True)
 class Commands(Layer):
     commands: Tuple[str, ...]
+    secret_mounts: Optional[Tuple[str | Secret, ...]] = None
 
     def update_hash(self, hasher: hashlib._Hash):
-        hasher.update("".join(self.commands).encode("utf-8"))
+        hash_input = "".join(self.commands)
+
+        if self.secret_mounts:
+            for secret_mount in self.secret_mounts:
+                hash_input += str(secret_mount)
+        hasher.update(hash_input.encode("utf-8"))
 
 
 @rich.repr.auto
@@ -279,7 +285,7 @@ class DockerIgnore(Layer):
 @dataclass(frozen=True, repr=True)
 class CopyConfig(Layer):
     path_type: CopyConfigType = field(metadata={"identifier": True})
-    src: Path = field(metadata={"identifier": True})
+    src: Path = field(metadata={"identifier": False})
     dst: str
     src_name: str = field(init=False)
 
@@ -451,11 +457,7 @@ class Image:
         # this default image definition may need to be updated once there is a released pypi version
         from flyte._version import __version__
 
-        dev_mode = (
-            (cls._is_editable_install() or (__version__ and "dev" in __version__))
-            and not flyte_version
-            and install_flyte
-        )
+        dev_mode = (__version__ and "dev" in __version__) and not flyte_version and install_flyte
         if install_flyte is False:
             preset_tag = f"py{python_version[0]}.{python_version[1]}"
         else:
@@ -507,13 +509,6 @@ class Image:
 
         return image
 
-    @staticmethod
-    def _is_editable_install():
-        """Internal hacky function to see if the current install is editable or not."""
-        curr = Path(__file__)
-        pyproject = curr.parent.parent.parent / "pyproject.toml"
-        return pyproject.exists()
-
     @classmethod
     def from_debian_base(
         cls,
@@ -793,9 +788,24 @@ class Image:
 
         Example:
         ```python
-        @flyte.task(image=(flyte.Image
-                        .ubuntu_python()
-                        .with_pip_packages("requests", "numpy")))
+        @flyte.task(image=(flyte.Image.from_debian_base().with_pip_packages("requests", "numpy")))
+        def my_task(x: int) -> int:
+            import numpy as np
+            return np.sum([x, 1])
+        ```
+
+        To mount secrets during the build process to download private packages, you can use the `secret_mounts`.
+        In the below example, "GITHUB_PAT" will be mounted as env var "GITHUB_PAT",
+         and "apt-secret" will be mounted at /etc/apt/apt-secret.
+        Example:
+        ```python
+        private_package = "git+https://$GITHUB_PAT@github.com/flyteorg/flytex.git@2e20a2acebfc3877d84af643fdd768edea41d533"
+        @flyte.task(
+            image=(
+                flyte.Image.from_debian_base()
+                .with_pip_packages("private_package", secret_mounts=[Secret(key="GITHUB_PAT")])
+                .with_apt_packages("git", secret_mounts=[Secret(key="apt-secret", mount="/etc/apt/apt-secret")])
+        )
         def my_task(x: int) -> int:
             import numpy as np
             return np.sum([x, 1])
@@ -920,16 +930,21 @@ class Image:
         )
         return new_image
 
-    def with_commands(self, commands: List[str]) -> Image:
+    def with_commands(self, commands: List[str], secret_mounts: Optional[SecretRequest] = None) -> Image:
         """
         Use this method to create a new image with the specified commands layered on top of the current image
         Be sure not to use RUN in your command.
 
         :param commands: list of commands to run
+        :param secret_mounts: list of secret mounts to use for the build process.
         :return: Image
         """
         new_commands: Tuple = _ensure_tuple(commands)
-        new_image = self.clone(addl_layer=Commands(commands=new_commands))
+        new_image = self.clone(
+            addl_layer=Commands(
+                commands=new_commands, secret_mounts=_ensure_tuple(secret_mounts) if secret_mounts else None
+            )
+        )
         return new_image
 
     def with_local_v2(self) -> Image:

flyte/_initialize.py

@@ -11,7 +11,6 @@ from flyte.errors import InitializationError
 from flyte.syncify import syncify
 
 from ._logging import initialize_logger, logger
-from ._tools import ipython_check
 
 if TYPE_CHECKING:
     from flyte._internal.imagebuild import ImageBuildEngine
@@ -173,6 +172,7 @@ async def init(
 
     :return: None
     """
+    from flyte._tools import ipython_check
     from flyte._utils import get_cwd_editable_install, org_from_endpoint, sanitize_endpoint
 
     interactive_mode = ipython_check()
@@ -205,6 +205,14 @@ async def init(
                 http_proxy_url=http_proxy_url,
             )
 
+        if not root_dir:
+            editable_root = get_cwd_editable_install()
+            if editable_root:
+                logger.info(f"Using editable install as root directory: {editable_root}")
+                root_dir = editable_root
+            else:
+                logger.info("No editable install found, using current working directory as root directory.")
+                root_dir = Path.cwd()
         root_dir = root_dir or get_cwd_editable_install() or Path.cwd()
         _init_config = _InitConfig(
             root_dir=root_dir,
@@ -242,17 +250,19 @@ async def init_from_config(
     cfg: config.Config
     if path_or_config is None or isinstance(path_or_config, str):
         # If a string is passed, treat it as a path to the config file
-        if path_or_config:
+        if root_dir and path_or_config:
+            cfg = config.auto(str(root_dir / path_or_config))
+        elif path_or_config:
             if not Path(path_or_config).exists():
                 raise InitializationError(
                     "ConfigFileNotFoundError",
                     "user",
                     f"Configuration file '{path_or_config}' does not exist., current working directory is {Path.cwd()}",
                 )
-        if root_dir and path_or_config:
-            cfg = config.auto(str(root_dir / path_or_config))
-        else:
             cfg = config.auto(path_or_config)
+        else:
+            # If no path is provided, use the default config file
+            cfg = config.auto()
     else:
         # If a Config object is passed, use it directly
         cfg = path_or_config
@@ -374,30 +384,6 @@ def ensure_client():
         )
 
 
-def requires_client(func: T) -> T:
-    """
-    Decorator that checks if the client has been initialized before executing the function.
-    Raises InitializationError if the client is not initialized.
-
-    :param func: Function to decorate
-    :return: Decorated function that checks for initialization
-    """
-
-    @functools.wraps(func)
-    async def wrapper(*args, **kwargs) -> T:
-        init_config = _get_init_config()
-        if init_config is None or init_config.client is None:
-            raise InitializationError(
-                "ClientNotInitializedError",
-                "user",
-                f"Function '{func.__name__}' requires client to be initialized. "
-                f"Call flyte.init() with a valid endpoint or api-key before using this function.",
-            )
-        return func(*args, **kwargs)
-
-    return typing.cast(T, wrapper)
-
-
 def requires_storage(func: T) -> T:
     """
     Decorator that checks if the storage has been initialized before executing the function.

flyte/_internal/controllers/remote/_action.py

@@ -183,7 +183,7 @@ class Action:
         et.FromSeconds(int(end_time))
         et.nanos = int((end_time % 1) * 1e9)
 
-        spec = (
+        _ = (
             task_definition_pb2.TaskSpec(task_template=tasks_pb2.TaskTemplate(interface=typed_interface))
             if typed_interface
             else None
@@ -208,6 +208,6 @@ class Action:
                     output_uri=outputs_uri,
                     report_uri=report_uri,
                 ),
-                spec=spec,
+                # spec=spec,
             ),
         )

flyte/_internal/controllers/remote/_controller.py

@@ -167,7 +167,7 @@ class RemoteController(Controller):
         # It is not allowed to change the code bundle (for regular code bundles) in the middle of a run.
         code_bundle = tctx.code_bundle
 
-        if code_bundle and code_bundle.pkl:
+        if tctx.interactive_mode or (code_bundle and code_bundle.pkl):
             logger.debug(f"Building new pkl bundle for task {_task.name}")
             code_bundle = await build_pkl_bundle(
                 _task,

flyte/_internal/imagebuild/docker_builder.py

@@ -302,11 +302,9 @@ def _get_secret_commands(layers: typing.Tuple[Layer, ...]) -> typing.List[str]:
     commands = []
 
     def _get_secret_command(secret: str | Secret) -> typing.List[str]:
-        secret_id = hash(secret)
         if isinstance(secret, str):
-            if not os.path.exists(secret):
-                raise FileNotFoundError(f"Secret file '{secret}' not found")
-            return ["--secret", f"id={secret_id},src={secret}"]
+            secret = Secret(key=secret)
+        secret_id = hash(secret)
         secret_env_key = "_".join([k.upper() for k in filter(None, (secret.group, secret.key))])
         secret_env = os.getenv(secret_env_key)
         if secret_env:
@@ -329,18 +327,16 @@ def _get_secret_mounts_layer(secrets: typing.Tuple[str | Secret, ...] | None) ->
     if secrets is None:
         return ""
     secret_mounts_layer = ""
-    for secret in secrets:
+    for s in secrets:
+        secret = Secret(key=s) if isinstance(s, str) else s
         secret_id = hash(secret)
-        if isinstance(secret, str):
-            secret_mounts_layer += f"--mount=type=secret,id={secret_id},target=/run/secrets/{os.path.basename(secret)}"
-        elif isinstance(secret, Secret):
-            if secret.mount:
-                secret_mounts_layer += f"--mount=type=secret,id={secret_id},target={secret.mount}"
-            elif secret.as_env_var:
-                secret_mounts_layer += f"--mount=type=secret,id={secret_id},env={secret.as_env_var}"
-            else:
-                secret_file_name = "_".join(list(filter(None, (secret.group, secret.key))))
-                secret_mounts_layer += f"--mount=type=secret,id={secret_id},src=/run/secrets/{secret_file_name}"
+        if secret.mount:
+            secret_mounts_layer += f"--mount=type=secret,id={secret_id},target={secret.mount}"
+        elif secret.as_env_var:
+            secret_mounts_layer += f"--mount=type=secret,id={secret_id},env={secret.as_env_var}"
+        else:
+            secret_default_env_key = "_".join(list(filter(None, (secret.group, secret.key))))
+            secret_mounts_layer += f"--mount=type=secret,id={secret_id},env={secret_default_env_key}"
 
     return secret_mounts_layer
 

flyte/_internal/imagebuild/remote_builder.py

@@ -19,21 +19,26 @@ from flyte._image import (
     CopyConfig,
     DockerIgnore,
     Env,
+    PipOption,
     PipPackages,
     PythonWheels,
     Requirements,
     UVProject,
     UVScript,
+    WorkDir,
 )
 from flyte._internal.imagebuild.image_builder import ImageBuilder, ImageChecker
 from flyte._internal.imagebuild.utils import copy_files_to_context
+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 flyte._protos.imagebuilder import definition_pb2 as image_definition_pb2
 
 IMAGE_TASK_NAME = "build-image"
+OPTIMIZE_TASK_NAME = "optimize-task"
 IMAGE_TASK_PROJECT = "system"
 IMAGE_TASK_DOMAIN = "production"
 
@@ -95,12 +100,12 @@ class RemoteImageBuilder(ImageBuilder):
         spec, context = await _validate_configuration(image)
 
         start = datetime.now(timezone.utc)
-        entity = remote.Task.get(
+        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))
         run = cast(
             Run,
             await flyte.with_runcontext(project=IMAGE_TASK_PROJECT, domain=IMAGE_TASK_DOMAIN).run.aio(
@@ -117,6 +122,19 @@ class RemoteImageBuilder(ImageBuilder):
 
         if run_details.action_details.raw_phase == run_definition_pb2.PHASE_SUCCEEDED:
             logger.warning(click.style(f"✅ Build completed in {elapsed}!", bold=True, fg="green"))
+            try:
+                entity = remote.Task.get(
+                    name=OPTIMIZE_TASK_NAME,
+                    project=IMAGE_TASK_PROJECT,
+                    domain=IMAGE_TASK_DOMAIN,
+                    auto_version="latest",
+                )
+                await flyte.with_runcontext(project=IMAGE_TASK_PROJECT, domain=IMAGE_TASK_DOMAIN).run.aio(
+                    entity, spec=spec, context=context, target_image=image_name
+                )
+            except Exception as e:
+                # Ignore the error if optimize is not enabled in the backend.
+                logger.warning(f"Failed to run optimize task with error: {e}")
         else:
             raise flyte.errors.ImageBuildError(f"❌ Build failed in {elapsed} at {click.style(run.url, fg='cyan')}")
 
@@ -164,9 +182,27 @@ def _get_layers_proto(image: Image, context_path: Path) -> "image_definition_pb2
 
     layers = []
     for layer in image._layers:
+        secret_mounts = None
+        pip_options = None
+
+        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)
+                apt_packages=image_definition_pb2.AptPackages(
+                    packages=layer.packages,
+                    secret_mounts=secret_mounts,
+                ),
             )
             layers.append(apt_layer)
         elif isinstance(layer, PythonWheels):
@@ -174,12 +210,8 @@ def _get_layers_proto(image: Image, context_path: Path) -> "image_definition_pb2
             wheel_layer = image_definition_pb2.Layer(
                 python_wheels=image_definition_pb2.PythonWheels(
                     dir=str(dst_path.relative_to(context_path)),
-                    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,
-                    ),
+                    options=pip_options,
+                    secret_mounts=secret_mounts,
                 )
             )
             layers.append(wheel_layer)
@@ -189,12 +221,8 @@ def _get_layers_proto(image: Image, context_path: Path) -> "image_definition_pb2
             requirements_layer = image_definition_pb2.Layer(
                 requirements=image_definition_pb2.Requirements(
                     file=str(dst_path.relative_to(context_path)),
-                    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,
-                    ),
+                    options=pip_options,
+                    secret_mounts=secret_mounts,
                 )
             )
             layers.append(requirements_layer)
@@ -211,12 +239,8 @@ def _get_layers_proto(image: Image, context_path: Path) -> "image_definition_pb2
             pip_layer = image_definition_pb2.Layer(
                 pip_packages=image_definition_pb2.PipPackages(
                     packages=packages,
-                    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,
-                    ),
+                    options=pip_options,
+                    secret_mounts=secret_mounts,
                 )
             )
             layers.append(pip_layer)
@@ -235,7 +259,10 @@ def _get_layers_proto(image: Image, context_path: Path) -> "image_definition_pb2
             layers.append(uv_layer)
         elif isinstance(layer, Commands):
             commands_layer = image_definition_pb2.Layer(
-                commands=image_definition_pb2.Commands(cmd=list(layer.commands))
+                commands=image_definition_pb2.Commands(
+                    cmd=list(layer.commands),
+                    secret_mounts=secret_mounts,
+                )
             )
             layers.append(commands_layer)
         elif isinstance(layer, DockerIgnore):
@@ -257,6 +284,11 @@ def _get_layers_proto(image: Image, context_path: Path) -> "image_definition_pb2
                 )
             )
             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,
@@ -267,3 +299,20 @@ def _get_layers_proto(image: Image, context_path: Path) -> "image_definition_pb2
 
 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)}")
+
+    return secrets

flyte/_internal/runtime/entrypoints.py

@@ -142,6 +142,7 @@ async def load_and_run_task(
     code_bundle: CodeBundle | None = None,
     input_path: str | None = None,
     image_cache: ImageCache | None = None,
+    interactive_mode: bool = False,
 ):
     """
     This method is invoked from the runtime/CLI and is used to run a task. This creates the context tree,
@@ -159,6 +160,7 @@ async def load_and_run_task(
     :param code_bundle: The code bundle to use for the task.
     :param input_path: The input path to use for the task.
     :param image_cache: Mappings of Image identifiers to image URIs.
+    :param interactive_mode: Whether to run the task in interactive mode.
     """
     task = await _download_and_load_task(code_bundle, resolver, resolver_args)
 
@@ -175,4 +177,5 @@ async def load_and_run_task(
         code_bundle=code_bundle,
         input_path=input_path,
         image_cache=image_cache,
+        interactive_mode=interactive_mode,
     )

flyte/_internal/runtime/reuse.py

@@ -54,10 +54,10 @@ def extract_unique_id_and_image(
         components += f":{reuse_policy.replicas}"
     if reuse_policy.ttl is not None:
         components += f":{reuse_policy.ttl.total_seconds()}"
-    if reuse_policy.reuse_salt is None and code_bundle is not None:
+    if reuse_policy.get_scaledown_ttl() is not None:
+        components += f":{reuse_policy.get_scaledown_ttl()}"
+    if code_bundle is not None:
         components += f":{code_bundle.computed_version}"
-    else:
-        components += f":{reuse_policy.reuse_salt}"
     if task.security_context is not None:
         security_ctx_str = task.security_context.SerializeToString(deterministic=True)
         components += f":{security_ctx_str}"
@@ -102,6 +102,8 @@ def add_reusable(
         env_name=name, code_bundle=code_bundle, task=task, reuse_policy=reuse_policy
     )
 
+    scaledown_ttl = reuse_policy.get_scaledown_ttl()
+
     task.custom = {
         "name": name,
         "version": version[:15],  # Use only the first 15 characters for the version
@@ -110,8 +112,10 @@ def add_reusable(
             "container_image": image_uri,
             "backlog_length": None,
             "parallelism": reuse_policy.concurrency,
+            "min_replica_count": reuse_policy.min_replicas,
             "replica_count": reuse_policy.max_replicas,
             "ttl_seconds": reuse_policy.ttl.total_seconds() if reuse_policy.ttl else None,
+            "scaledown_ttl_seconds": scaledown_ttl.total_seconds() if scaledown_ttl else None,
         },
     }
 

flyte/_internal/runtime/task_serde.py

@@ -54,7 +54,7 @@ def translate_task_to_wire(
     return task_definition_pb2.TaskSpec(
         task_template=tt,
         default_inputs=default_inputs,
-        short_name=task.friendly_name[:_MAX_TASK_SHORT_NAME_LENGTH],
+        short_name=task.short_name[:_MAX_TASK_SHORT_NAME_LENGTH],
         environment=env,
     )
 
@@ -145,7 +145,6 @@ def get_proto_task(task: TaskTemplate, serialize_context: SerializationContext)
             logger.debug(f"Detected pkl bundle for task {task.name}, using computed version as cache version")
             cache_version = serialize_context.code_bundle.computed_version
         else:
-            version_parameters = None
             if isinstance(task, AsyncFunctionTaskTemplate):
                 version_parameters = VersionParameters(func=task.func, image=task.image)
             else:
@@ -204,7 +203,9 @@ def _get_urun_container(
     serialize_context: SerializationContext, task_template: TaskTemplate
 ) -> Optional[tasks_pb2.Container]:
     env = (
-        [literals_pb2.KeyValuePair(key=k, value=v) for k, v in task_template.env.items()] if task_template.env else None
+        [literals_pb2.KeyValuePair(key=k, value=v) for k, v in task_template.env_vars.items()]
+        if task_template.env_vars
+        else None
     )
     resources = get_proto_resources(task_template.resources)
     # pr: under what conditions should this return None?

flyte/_internal/runtime/taskrunner.py

@@ -110,16 +110,18 @@ async def run_task(
 async def convert_and_run(
     *,
     task: TaskTemplate,
-    inputs: Inputs,
     action: ActionID,
     controller: Controller,
     raw_data_path: RawDataPath,
     version: str,
     output_path: str,
     run_base_dir: str,
+    inputs: Inputs = Inputs.empty(),
+    input_path: str | None = None,
     checkpoints: Checkpoints | None = None,
     code_bundle: CodeBundle | None = None,
     image_cache: ImageCache | None = None,
+    interactive_mode: bool = False,
 ) -> Tuple[Optional[Outputs], Optional[Error]]:
     """
     This method is used to convert the inputs to native types, and run the task. It assumes you are running
@@ -130,6 +132,7 @@ async def convert_and_run(
         action=action,
         checkpoints=checkpoints,
         code_bundle=code_bundle,
+        input_path=input_path,
         output_path=output_path,
         run_base_dir=run_base_dir,
         version=version,
@@ -137,8 +140,10 @@ async def convert_and_run(
         compiled_image_cache=image_cache,
         report=flyte.report.Report(name=action.name),
         mode="remote" if not ctx.data.task_context else ctx.data.task_context.mode,
+        interactive_mode=interactive_mode,
     )
     with ctx.replace_task_context(tctx):
+        inputs = await load_inputs(input_path) if input_path else inputs
         inputs_kwargs = await convert_inputs_to_native(inputs, task.native_interface)
         out, err = await run_task(tctx=tctx, controller=controller, task=task, inputs=inputs_kwargs)
         if err is not None:
@@ -161,15 +166,15 @@ async def extract_download_run_upload(
     code_bundle: CodeBundle | None = None,
     input_path: str | None = None,
     image_cache: ImageCache | None = None,
+    interactive_mode: bool = False,
 ):
     """
     This method is invoked from the CLI (urun) and is used to run a task. This assumes that the context tree
     has already been created, and the task has been loaded. It also handles the loading of the task.
     """
-    inputs = await load_inputs(input_path) if input_path else None
     outputs, err = await convert_and_run(
         task=task,
-        inputs=inputs or Inputs.empty(),
+        input_path=input_path,
         action=action,
         controller=controller,
         raw_data_path=raw_data_path,
@@ -179,6 +184,7 @@ async def extract_download_run_upload(
         checkpoints=checkpoints,
         code_bundle=code_bundle,
         image_cache=image_cache,
+        interactive_mode=interactive_mode,
     )
     if err is not None:
         path = await upload_error(err.err, output_path)

flyte/_logging.py

@@ -4,7 +4,9 @@ import logging
 import os
 from typing import Optional
 
-from ._tools import ipython_check, is_in_cluster
+import flyte
+
+from ._tools import ipython_check
 
 DEFAULT_LOG_LEVEL = logging.WARNING
 
@@ -42,7 +44,8 @@ def get_rich_handler(log_level: int) -> Optional[logging.Handler]:
     """
     Upgrades the global loggers to use Rich logging.
     """
-    if is_in_cluster():
+    ctx = flyte.ctx()
+    if ctx and ctx.is_in_cluster():
         return None
     if not ipython_check() and is_rich_logging_disabled():
         return None