flyte 2.0.0b8__py3-none-any.whl → 2.0.0b13__py3-none-any.whl

flyte/__init__.py

@@ -2,37 +2,7 @@
 Flyte SDK for authoring compound AI applications, services and workflows.
 """
 
-__all__ = [
-    "GPU",
-    "TPU",
-    "Cache",
-    "CachePolicy",
-    "CacheRequest",
-    "Device",
-    "Environment",
-    "Image",
-    "PodTemplate",
-    "Resources",
-    "RetryStrategy",
-    "ReusePolicy",
-    "Secret",
-    "SecretRequest",
-    "TaskEnvironment",
-    "Timeout",
-    "TimeoutType",
-    "__version__",
-    "build",
-    "build_images",
-    "ctx",
-    "deploy",
-    "group",
-    "init",
-    "init_from_config",
-    "map",
-    "run",
-    "trace",
-    "with_runcontext",
-]
+from __future__ import annotations
 
 import sys
 
@@ -60,8 +30,62 @@ from ._version import __version__
 sys.excepthook = custom_excepthook
 
 
+def _silence_grpc_warnings():
+    """
+    Silences gRPC warnings that can clutter the output.
+    """
+    import os
+
+    # Set environment variables for gRPC, this reduces log spew and avoids unnecessary warnings
+    # before importing grpc
+    if "GRPC_VERBOSITY" not in os.environ:
+        os.environ["GRPC_VERBOSITY"] = "ERROR"
+        os.environ["GRPC_CPP_MIN_LOG_LEVEL"] = "ERROR"
+        # Disable fork support (stops "skipping fork() handlers")
+        os.environ["GRPC_ENABLE_FORK_SUPPORT"] = "0"
+        # Reduce absl/glog verbosity
+        os.environ["GLOG_minloglevel"] = "2"
+        os.environ["ABSL_LOG"] = "0"
+
+
+_silence_grpc_warnings()
+
+
 def version() -> str:
     """
     Returns the version of the Flyte SDK.
     """
     return __version__
+
+
+__all__ = [
+    "GPU",
+    "TPU",
+    "Cache",
+    "CachePolicy",
+    "CacheRequest",
+    "Device",
+    "Environment",
+    "Image",
+    "PodTemplate",
+    "Resources",
+    "RetryStrategy",
+    "ReusePolicy",
+    "Secret",
+    "SecretRequest",
+    "TaskEnvironment",
+    "Timeout",
+    "TimeoutType",
+    "__version__",
+    "build",
+    "build_images",
+    "ctx",
+    "deploy",
+    "group",
+    "init",
+    "init_from_config",
+    "map",
+    "run",
+    "trace",
+    "with_runcontext",
+]

flyte/_context.py

@@ -97,7 +97,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
 

flyte/_environment.py

@@ -36,7 +36,7 @@ class Environment:
     :param name: Name of the environment
     :param image: Docker image to use for the environment. If set to "auto", will use the default image.
     :param resources: Resources to allocate for the environment.
-    :param env: Environment variables to set for the environment.
+    :param env_vars: Environment variables to set for the environment.
     :param secrets: Secrets to inject into the environment.
     :param depends_on: Environment dependencies to hint, so when you deploy the environment, the dependencies are
         also deployed. This is useful when you have a set of environments that depend on each other.
@@ -47,7 +47,7 @@ class Environment:
     pod_template: Optional[Union[str, "V1PodTemplate"]] = None
     description: Optional[str] = None
     secrets: Optional[SecretRequest] = None
-    env: Optional[Dict[str, str]] = None
+    env_vars: Optional[Dict[str, str]] = None
     resources: Optional[Resources] = None
     image: Union[str, Image, Literal["auto"]] = "auto"
 
@@ -75,7 +75,7 @@ class Environment:
         name: str,
         image: Optional[Union[str, Image, Literal["auto"]]] = None,
         resources: Optional[Resources] = None,
-        env: Optional[Dict[str, str]] = None,
+        env_vars: Optional[Dict[str, str]] = None,
         secrets: Optional[SecretRequest] = None,
         depends_on: Optional[List[Environment]] = None,
         **kwargs: Any,
@@ -94,8 +94,8 @@ class Environment:
             kwargs["resources"] = self.resources
         if self.secrets is not None:
             kwargs["secrets"] = self.secrets
-        if self.env is not None:
-            kwargs["env"] = self.env
+        if self.env_vars is not None:
+            kwargs["env_vars"] = self.env_vars
         if self.pod_template is not None:
             kwargs["pod_template"] = self.pod_template
         if self.description is not None:

flyte/_image.py

@@ -279,7 +279,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 +451,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 +503,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,

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()

flyte/_internal/imagebuild/docker_builder.py

@@ -61,7 +61,7 @@ ENV PATH="/root/.venv/bin:$$PATH" \
 
 UV_PACKAGE_INSTALL_COMMAND_TEMPLATE = Template("""\
 RUN --mount=type=cache,sharing=locked,mode=0777,target=/root/.cache/uv,id=uv \
-    --mount=type=bind,target=requirements_uv.txt,src=requirements_uv.txt \
+    $REQUIREMENTS_MOUNT \
     $SECRET_MOUNT \
     uv pip install --python $$UV_PYTHON $PIP_INSTALL_ARGS
 """)
@@ -129,6 +129,8 @@ class Handler(Protocol):
 class PipAndRequirementsHandler:
     @staticmethod
     async def handle(layer: PipPackages, context_path: Path, dockerfile: str) -> str:
+        secret_mounts = _get_secret_mounts_layer(layer.secret_mounts)
+
         # Set pip_install_args based on the layer type - either a requirements file or a list of packages
         if isinstance(layer, Requirements):
             if not layer.file.exists():
@@ -138,17 +140,20 @@ class PipAndRequirementsHandler:
 
             # Copy the requirements file to the context path
             requirements_path = copy_files_to_context(layer.file, context_path)
+            rel_path = str(requirements_path.relative_to(context_path))
             pip_install_args = layer.get_pip_install_args()
-            pip_install_args.extend(["--requirement", str(requirements_path)])
+            pip_install_args.extend(["--requirement", "requirements.txt"])
+            mount = f"--mount=type=bind,target=requirements.txt,src={rel_path}"
         else:
+            mount = ""
             requirements = list(layer.packages) if layer.packages else []
             reqs = " ".join(requirements)
             pip_install_args = layer.get_pip_install_args()
             pip_install_args.append(reqs)
 
-        secret_mounts = _get_secret_mounts_layer(layer.secret_mounts)
         delta = UV_PACKAGE_INSTALL_COMMAND_TEMPLATE.substitute(
             SECRET_MOUNT=secret_mounts,
+            REQUIREMENTS_MOUNT=mount,
             PIP_INSTALL_ARGS=" ".join(pip_install_args),
         )
 

flyte/_internal/imagebuild/remote_builder.py

@@ -24,6 +24,7 @@ from flyte._image import (
     Requirements,
     UVProject,
     UVScript,
+    WorkDir,
 )
 from flyte._internal.imagebuild.image_builder import ImageBuilder, ImageChecker
 from flyte._internal.imagebuild.utils import copy_files_to_context
@@ -34,6 +35,7 @@ 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"
 
@@ -117,6 +119,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')}")
 
@@ -257,6 +272,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,

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

@@ -204,7 +204,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/_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

flyte/_reusable_environment.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 from datetime import timedelta
-from typing import Optional, Tuple, Union
+from typing import Tuple, Union
 
 from flyte._logging import logger
 
@@ -17,25 +17,22 @@ class ReusePolicy:
 
     :param replicas: Either a single int representing number of replicas or a tuple of two ints representing
      the min and max.
-    :param idle_ttl: The maximum idle duration for an environment replica, specified as either seconds (int) or a
-        timedelta. If not set, the environment's global default will be used.
+    :param idle_ttl: The maximum idle duration for an environment, specified as either seconds (int) or a
+        timedelta, after which all replicas in the environment are shutdown.
+        If not set, the default is configured in the backend (can be as low as 90s).
         When a replica remains idle — meaning no tasks are running — for this duration, it will be automatically
-        terminated.
+        terminated, also referred to as environment idle timeout.
     :param concurrency: The maximum number of tasks that can run concurrently in one instance of the environment.
-          Concurrency of greater than 1 is only supported only for `async` tasks.
-    :param reuse_salt: Optional string used to control environment reuse.
-        If set, the environment will be reused even if the code bundle changes.
-        To force a new environment, either set this to `None` or change its value.
-
-        Example:
-            reuse_salt = "v1"  # Environment is reused
-            reuse_salt = "v2"  # Forces environment recreation
+          Concurrency of greater than 1 is only supported for `async` tasks.
+    :param scaledown_ttl: The minimum time to wait before scaling down each replica, specified as either seconds (int)
+        or a timedelta. This is useful to prevent rapid scaling down of replicas when tasks are running
+        frequently. If not set, the default is configured in the backend.
     """
 
     replicas: Union[int, Tuple[int, int]] = 2
-    idle_ttl: Optional[Union[int, timedelta]] = None
-    reuse_salt: str | None = None
+    idle_ttl: Union[int, timedelta] = 30  # seconds
     concurrency: int = 1
+    scaledown_ttl: Union[int, timedelta] = 30  # seconds
 
     def __post_init__(self):
         if self.replicas is None:
@@ -47,11 +44,12 @@ class ReusePolicy:
         elif len(self.replicas) != 2:
             raise ValueError("replicas must be an int or a tuple of two ints")
 
-        if self.idle_ttl:
-            if isinstance(self.idle_ttl, int):
-                self.idle_ttl = timedelta(seconds=int(self.idle_ttl))
-            elif not isinstance(self.idle_ttl, timedelta):
-                raise ValueError("idle_ttl must be an int (seconds) or a timedelta")
+        if isinstance(self.idle_ttl, int):
+            self.idle_ttl = timedelta(seconds=int(self.idle_ttl))
+        elif not isinstance(self.idle_ttl, timedelta):
+            raise ValueError("idle_ttl must be an int (seconds) or a timedelta")
+        if self.idle_ttl.total_seconds() < 30:
+            raise ValueError("idle_ttl must be at least 30 seconds")
 
         if self.replicas[1] == 1 and self.concurrency == 1:
             logger.warning(
@@ -61,6 +59,13 @@ class ReusePolicy:
                 "that runs child tasks."
             )
 
+        if isinstance(self.scaledown_ttl, int):
+            self.scaledown_ttl = timedelta(seconds=int(self.scaledown_ttl))
+        elif not isinstance(self.scaledown_ttl, timedelta):
+            raise ValueError("scaledown_ttl must be an int (seconds) or a timedelta")
+        if self.scaledown_ttl.total_seconds() < 30:
+            raise ValueError("scaledown_ttl must be at least 30 seconds")
+
     @property
     def ttl(self) -> timedelta | None:
         """
@@ -72,6 +77,23 @@ class ReusePolicy:
             return self.idle_ttl
         return timedelta(seconds=self.idle_ttl)
 
+    @property
+    def min_replicas(self) -> int:
+        """
+        Returns the minimum number of replicas.
+        """
+        return self.replicas[0] if isinstance(self.replicas, tuple) else self.replicas
+
+    def get_scaledown_ttl(self) -> timedelta | None:
+        """
+        Returns the scaledown TTL as a timedelta. If scaledown_ttl is not set, returns None.
+        """
+        if self.scaledown_ttl is None:
+            return None
+        if isinstance(self.scaledown_ttl, timedelta):
+            return self.scaledown_ttl
+        return timedelta(seconds=int(self.scaledown_ttl))
+
     @property
     def max_replicas(self) -> int:
         """

flyte/_run.py

@@ -19,8 +19,6 @@ from flyte._initialize import (
 )
 from flyte._logging import logger
 from flyte._task import P, R, TaskTemplate
-from flyte._tools import ipython_check
-from flyte.errors import InitializationError
 from flyte.models import (
     ActionID,
     Checkpoints,
@@ -89,13 +87,15 @@ class _Runner:
         overwrite_cache: bool = False,
         project: str | None = None,
         domain: str | None = None,
-        env: Dict[str, str] | None = None,
+        env_vars: Dict[str, str] | None = None,
         labels: Dict[str, str] | None = None,
         annotations: Dict[str, str] | None = None,
         interruptible: bool = False,
         log_level: int | None = None,
         disable_run_cache: bool = False,
     ):
+        from flyte._tools import ipython_check
+
         init_config = _get_init_config()
         client = init_config.client if init_config else None
         if not force_mode and client is not None:
@@ -116,7 +116,7 @@ class _Runner:
         self._overwrite_cache = overwrite_cache
         self._project = project
         self._domain = domain
-        self._env = env
+        self._env_vars = env_vars
         self._labels = labels
         self._annotations = annotations
         self._interruptible = interruptible
@@ -198,7 +198,7 @@ class _Runner:
             task_spec = translate_task_to_wire(obj, s_ctx)
             inputs = await convert_from_native_to_inputs(obj.native_interface, *args, **kwargs)
 
-        env = self._env or {}
+        env = self._env_vars or {}
         if self._log_level:
             env["LOG_LEVEL"] = str(self._log_level)
         else:
@@ -207,7 +207,7 @@ class _Runner:
         if not self._dry_run:
             if get_client() is None:
                 # This can only happen, if the user forces flyte.run(mode="remote") without initializing the client
-                raise InitializationError(
+                raise flyte.errors.InitializationError(
                     "ClientNotInitializedError",
                     "user",
                     "flyte.run requires client to be initialized. "
@@ -542,7 +542,7 @@ def with_runcontext(
     overwrite_cache: bool = False,
     project: str | None = None,
     domain: str | None = None,
-    env: Dict[str, str] | None = None,
+    env_vars: Dict[str, str] | None = None,
     labels: Dict[str, str] | None = None,
     annotations: Dict[str, str] | None = None,
     interruptible: bool = False,
@@ -582,7 +582,7 @@ def with_runcontext(
     :param overwrite_cache: Optional If true, the cache will be overwritten for the run
     :param project: Optional The project to use for the run
     :param domain: Optional The domain to use for the run
-    :param env: Optional Environment variables to set for the run
+    :param env_vars: Optional Environment variables to set for the run
     :param labels: Optional Labels to set for the run
     :param annotations: Optional Annotations to set for the run
     :param interruptible: Optional If true, the run can be interrupted by the user.
@@ -606,7 +606,7 @@ def with_runcontext(
         raw_data_path=raw_data_path,
         run_base_dir=run_base_dir,
         overwrite_cache=overwrite_cache,
-        env=env,
+        env_vars=env_vars,
         labels=labels,
         annotations=annotations,
         interruptible=interruptible,

flyte/_secret.py

@@ -17,14 +17,14 @@ class Secret:
 
     Example:
     ```python
-    @task(secrets="MY_SECRET")
+    @task(secrets="my-secret")
     async def my_task():
-        os.environ["MY_SECRET"]  # This will be set to the value of the secret
+        # This will be set to the value of the secret. Note: The env var is always uppercase, and - is replaced with _.
+        os.environ["MY_SECRET"]
 
-    @task(secrets=Secret("MY_SECRET", mount="/path/to/secret"))
+    @task(secrets=Secret("my-openai-api-key", as_env_var="OPENAI_API_KEY"))
     async def my_task2():
-        async with open("/path/to/secret") as f:
-            secret_value = f.read()
+        os.environ["OPENAI_API_KEY"]
     ```
 
     TODO: Add support for secret versioning (some stores) and secret groups (some stores) and mounting as files.
@@ -32,6 +32,7 @@ class Secret:
     :param key: The name of the secret in the secret store.
     :param group: The group of the secret in the secret store.
     :param mount: Use this to specify the path where the secret should be mounted.
+    TODO: support arbitrary mount paths. Today only "/etc/flyte/secrets" is supported
     :param as_env_var: The name of the environment variable that the secret should be mounted as.
     """
 
@@ -41,6 +42,9 @@ class Secret:
     as_env_var: Optional[str] = None
 
     def __post_init__(self):
+        if not self.mount and not self.as_env_var:
+            self.as_env_var = f"{self.group}_{self.key}" if self.group else self.key
+            self.as_env_var = self.as_env_var.replace("-", "_").upper()
         if self.as_env_var is not None:
             pattern = r"^[A-Z_][A-Z0-9_]*$"
             if not re.match(pattern, self.as_env_var):

flyte/_task.py

@@ -76,7 +76,7 @@ class TaskTemplate(Generic[P, R]):
     :param reusable: Optional The reusability policy for the task, defaults to None, which means the task environment
     will not be reused across task invocations.
     :param docs: Optional The documentation for the task, if not provided the function docstring will be used.
-    :param env: Optional The environment variables to set for the task.
+    :param env_vars: Optional The environment variables to set for the task.
     :param secrets: Optional The secrets that will be injected into the task at runtime.
     :param timeout: Optional The timeout for the task.
     :param max_inline_io_bytes: Maximum allowed size (in bytes) for all inputs and outputs passed directly to the task
@@ -95,7 +95,7 @@ class TaskTemplate(Generic[P, R]):
     retries: Union[int, RetryStrategy] = 0
     reusable: Union[ReusePolicy, None] = None
     docs: Optional[Documentation] = None
-    env: Optional[Dict[str, str]] = None
+    env_vars: Optional[Dict[str, str]] = None
     secrets: Optional[SecretRequest] = None
     timeout: Optional[TimeoutType] = None
     pod_template: Optional[Union[str, PodTemplate]] = None
@@ -314,14 +314,16 @@ class TaskTemplate(Generic[P, R]):
     def override(
         self,
         *,
+        friendly_name: Optional[str] = None,
         resources: Optional[Resources] = None,
-        cache: CacheRequest = "auto",
+        cache: Optional[CacheRequest] = None,
         retries: Union[int, RetryStrategy] = 0,
         timeout: Optional[TimeoutType] = None,
         reusable: Union[ReusePolicy, Literal["off"], None] = None,
-        env: Optional[Dict[str, str]] = None,
+        env_vars: Optional[Dict[str, str]] = None,
         secrets: Optional[SecretRequest] = None,
         max_inline_io_bytes: int | None = None,
+        pod_template: Optional[Union[str, PodTemplate]] = None,
         **kwargs: Any,
     ) -> TaskTemplate:
         """
@@ -344,7 +346,7 @@ class TaskTemplate(Generic[P, R]):
                     " Reusable tasks will use the parent env's resources. You can disable reusability and"
                     " override resources if needed. (set reusable='off')"
                 )
-            if env is not None:
+            if env_vars is not None:
                 raise ValueError(
                     "Cannot override env when reusable is set."
                     " Reusable tasks will use the parent env's env. You can disable reusability and "
@@ -358,7 +360,7 @@ class TaskTemplate(Generic[P, R]):
                 )
 
         resources = resources or self.resources
-        env = env or self.env
+        env_vars = env_vars or self.env_vars
         secrets = secrets or self.secrets
 
         for k, v in kwargs.items():
@@ -373,14 +375,17 @@ class TaskTemplate(Generic[P, R]):
 
         return replace(
             self,
+            friendly_name=friendly_name or self.friendly_name,
             resources=resources,
             cache=cache,
             retries=retries,
             timeout=timeout,
             reusable=cast(Optional[ReusePolicy], reusable),
-            env=env,
+            env_vars=env_vars,
             secrets=secrets,
             max_inline_io_bytes=max_inline_io_bytes,
+            pod_template=pod_template,
+            **kwargs,
         )