flyte 0.2.0b9__py3-none-any.whl → 0.2.0b11__py3-none-any.whl

flyte/__init__.py

@@ -25,6 +25,7 @@ __all__ = [
     "Device",
     "Environment",
     "Image",
+    "PodTemplate",
     "Resources",
     "RetryStrategy",
     "ReusePolicy",
@@ -38,7 +39,7 @@ __all__ = [
     "deploy",
     "group",
     "init",
-    "init_auto_from_config",
+    "init_from_config",
     "map",
     "run",
     "trace",
@@ -51,8 +52,9 @@ from ._deploy import deploy
 from ._environment import Environment
 from ._group import group
 from ._image import Image
-from ._initialize import init, init_auto_from_config
+from ._initialize import init, init_from_config
 from ._map import map
+from ._pod import PodTemplate
 from ._resources import GPU, TPU, Device, Resources
 from ._retry import RetryStrategy
 from ._reusable_environment import ReusePolicy

flyte/_bin/runtime.py

@@ -76,13 +76,17 @@ def main(
 ):
     sys.path.insert(0, ".")
 
+    import flyte
     import flyte._utils as utils
     from flyte._initialize import initialize_in_cluster
     from flyte._internal.controllers import create_controller
     from flyte._internal.imagebuild.image_builder import ImageCache
     from flyte._internal.runtime.entrypoints import load_and_run_task
+    from flyte._logging import logger
     from flyte.models import ActionID, Checkpoints, CodeBundle, RawDataPath
 
+    logger.info(f"Initializing flyte runtime - version {flyte.__version__}")
+
     assert org, "Org is required for now"
     assert project, "Project is required"
     assert domain, "Domain is required"
@@ -98,15 +102,14 @@ def main(
     # This detection of api key is a hack for now.
     controller_kwargs: dict[str, Any] = {"insecure": False}
     if api_key := os.getenv(_UNION_EAGER_API_KEY_ENV_VAR):
-        from flyte._logging import logger
-
-        logger.warning(f"Using api key {api_key}")
+        logger.info("Using api key from environment")
         controller_kwargs["api_key"] = api_key
     else:
         ep = os.environ.get(ENDPOINT_OVERRIDE, "host.docker.internal:8090")
         controller_kwargs["endpoint"] = ep
         if "localhost" in ep or "docker" in ep:
             controller_kwargs["insecure"] = True
+        logger.debug(f"Using controller endpoint: {ep} with kwargs: {controller_kwargs}")
 
     bundle = CodeBundle(tgz=tgz, pkl=pkl, destination=dest, computed_version=version)
     initialize_in_cluster()

flyte/_deploy.py

@@ -128,6 +128,9 @@ async def apply(deployment: DeploymentPlan, copy_style: CopyFiles, dryrun: bool
     else:
         code_bundle = await build_code_bundle(from_dir=cfg.root_dir, dryrun=dryrun, copy_style=copy_style)
         deployment.version = code_bundle.computed_version
+        # TODO we should update the version to include the image cache digest and code bundle digest. This is
+        # to ensure that changes in image dependencies, cause an update to the deployment version.
+        # TODO Also hash the environment and tasks to ensure that changes in the environment or tasks
 
     sc = SerializationContext(
         project=cfg.project,

flyte/_initialize.py

@@ -138,14 +138,13 @@ async def init(
     :param project: Optional project name (not used in this implementation)
     :param domain: Optional domain name (not used in this implementation)
     :param root_dir: Optional root directory from which to determine how to load files, and find paths to files.
+      This is useful for determining the root directory for the current project, and for locating files like config etc.
+      also use to determine all the code that needs to be copied to the remote location.
       defaults to the editable install directory if the cwd is in a Python editable install, else just the cwd.
     :param log_level: Optional logging level for the logger, default is set using the default initialization policies
     :param api_key: Optional API key for authentication
     :param endpoint: Optional API endpoint URL
     :param headless: Optional Whether to run in headless mode
-    :param mode: Optional execution model (local, remote). Default is local. When local is used,
-           the execution will be done locally. When remote is used, the execution will be sent to a remote server,
-           In the remote case, the endpoint or api_key must be set.
     :param insecure_skip_verify: Whether to skip SSL certificate verification
     :param auth_client_config: Optional client configuration for authentication
     :param auth_type: The authentication type to use (Pkce, ClientSecret, ExternalCommand, DeviceFlow)
@@ -169,6 +168,8 @@ async def init(
     """
     from flyte._utils import get_cwd_editable_install
 
+    from ._utils.org_discovery import org_from_endpoint
+
     interactive_mode = ipython_check()
 
     initialize_logger(enable_rich=interactive_mode)
@@ -177,6 +178,9 @@ async def init(
 
     global _init_config  # noqa: PLW0603
 
+    if endpoint and "://" not in endpoint:
+        endpoint = f"dns:///{endpoint}"
+
     with _init_lock:
         client = None
         if endpoint or api_key:
@@ -204,17 +208,25 @@ async def init(
             domain=domain,
             client=client,
             storage=storage,
-            org=org,
+            org=org or org_from_endpoint(endpoint),
         )
 
 
 @syncify
-async def init_auto_from_config(path_or_config: str | Config | None = None) -> None:
+async def init_from_config(
+    path_or_config: str | Config | None = None, root_dir: Path | None = None, log_level: int | None = None
+) -> None:
     """
     Initialize the Flyte system using a configuration file or Config object. This method should be called before any
     other Flyte remote API methods are called. Thread-safe implementation.
 
     :param path_or_config: Path to the configuration file or Config object
+    :param root_dir: Optional root directory from which to determine how to load files, and find paths to
+        files like config etc. For example if one uses the copy-style=="all", it is essential to determine the
+        root directory for the current project. If not provided, it defaults to the editable install directory or
+        if not available, the current working directory.
+    :param log_level: Optional logging level for the framework logger,
+        default is set using the default initialization policies
     :return: None
     """
     import flyte.config as config
@@ -222,7 +234,17 @@ async def init_auto_from_config(path_or_config: str | Config | None = None) -> N
     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
-        cfg = config.auto(path_or_config)
+        if 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 a Config object is passed, use it directly
         cfg = path_or_config
@@ -241,6 +263,8 @@ async def init_auto_from_config(path_or_config: str | Config | None = None) -> N
         proxy_command=cfg.platform.proxy_command,
         client_id=cfg.platform.client_id,
         client_credentials_secret=cfg.platform.client_credentials_secret,
+        root_dir=root_dir,
+        log_level=log_level,
     )
 
 

flyte/_internal/controllers/_local_controller.py

@@ -14,7 +14,7 @@ from flyte._logging import log, logger
 from flyte._protos.workflow import task_definition_pb2
 from flyte._task import TaskTemplate
 from flyte._utils.helpers import _selector_policy
-from flyte.models import ActionID, NativeInterface, RawDataPath
+from flyte.models import ActionID, NativeInterface
 
 R = TypeVar("R")
 
@@ -86,7 +86,7 @@ class LocalController:
         sub_action_id, sub_action_output_path = convert.generate_sub_action_id_and_output_path(
             tctx, _task.name, serialized_inputs, 0
         )
-        sub_action_raw_data_path = RawDataPath(path=sub_action_output_path)
+        sub_action_raw_data_path = tctx.raw_data_path
 
         out, err = await direct_dispatch(
             _task,
@@ -162,6 +162,7 @@ class LocalController:
                 action=action_id,
                 interface=_interface,
                 inputs_path=action_output_path,
+                name=_func.__name__,
             ),
             True,
         )
@@ -179,7 +180,7 @@ class LocalController:
 
         if info.interface.outputs and info.output:
             # If the result is not an AsyncGenerator, convert it directly
-            converted_outputs = await convert.convert_from_native_to_outputs(info.output, info.interface)
+            converted_outputs = await convert.convert_from_native_to_outputs(info.output, info.interface, info.name)
             assert converted_outputs
         elif info.error:
             # If there is an error, convert it to a native error

flyte/_internal/controllers/_trace.py

@@ -18,6 +18,7 @@ class TraceInfo:
     duration: Optional[timedelta] = None
     output: Optional[Any] = None
     error: Optional[Exception] = None
+    name: str = ""
 
     def add_outputs(self, output: Any, duration: timedelta):
         """

flyte/_internal/controllers/remote/_action.py

@@ -130,7 +130,7 @@ class Action:
         """
         from flyte._logging import logger
 
-        logger.info(f"In Action from_state {obj.action_id} {obj.phase} {obj.output_uri}")
+        logger.debug(f"In Action from_state {obj.action_id} {obj.phase} {obj.output_uri}")
         return cls(
             action_id=obj.action_id,
             parent_action_name=parent_action_name,

flyte/_internal/controllers/remote/_informer.py

@@ -235,7 +235,7 @@ class Informer:
                     await self._shared_queue.put(node)
                     # hack to work in the absence of sentinel
             except asyncio.CancelledError:
-                logger.warning(f"Watch cancelled: {self.name}")
+                logger.info(f"Watch cancelled: {self.name}")
                 return
             except asyncio.TimeoutError as e:
                 logger.error(f"Watch timeout: {self.name}", exc_info=e)

flyte/_internal/runtime/convert.py

@@ -11,7 +11,7 @@ import flyte.errors
 import flyte.storage as storage
 from flyte._protos.workflow import run_definition_pb2, task_definition_pb2
 from flyte.models import ActionID, NativeInterface, TaskContext
-from flyte.types import TypeEngine
+from flyte.types import TypeEngine, TypeTransformerFailedError
 
 
 @dataclass(frozen=True)
@@ -80,7 +80,7 @@ async def convert_from_native_to_inputs(interface: NativeInterface, *args, **kwa
     )
 
 
-async def convert_from_native_to_outputs(o: Any, interface: NativeInterface) -> Outputs:
+async def convert_from_native_to_outputs(o: Any, interface: NativeInterface, task_name: str = "") -> Outputs:
     # Always make it a tuple even if it's just one item to simplify logic below
     if not isinstance(o, tuple):
         o = (o,)
@@ -90,8 +90,11 @@ async def convert_from_native_to_outputs(o: Any, interface: NativeInterface) ->
     )
     named = []
     for (output_name, python_type), v in zip(interface.outputs.items(), o):
-        lit = await TypeEngine.to_literal(v, python_type, TypeEngine.to_literal_type(python_type))
-        named.append(run_definition_pb2.NamedLiteral(name=output_name, value=lit))
+        try:
+            lit = await TypeEngine.to_literal(v, python_type, TypeEngine.to_literal_type(python_type))
+            named.append(run_definition_pb2.NamedLiteral(name=output_name, value=lit))
+        except TypeTransformerFailedError as e:
+            raise flyte.errors.RuntimeDataValidationError(output_name, e, task_name)
 
     return Outputs(proto_outputs=run_definition_pb2.Outputs(literals=named))
 

flyte/_internal/runtime/task_serde.py

@@ -3,9 +3,11 @@ This module provides functionality to serialize and deserialize tasks to and fro
 It includes a Resolver interface for loading tasks, and functions to load classes and tasks.
 """
 
+import copy
 import importlib
+import typing
 from datetime import timedelta
-from typing import Optional, Type
+from typing import Optional, Type, cast
 
 from flyteidl.core import identifier_pb2, literals_pb2, security_pb2, tasks_pb2
 from google.protobuf import duration_pb2, wrappers_pb2
@@ -13,6 +15,7 @@ from google.protobuf import duration_pb2, wrappers_pb2
 import flyte.errors
 from flyte._cache.cache import VersionParameters, cache_from_request
 from flyte._logging import logger
+from flyte._pod import _PRIMARY_CONTAINER_NAME_FIELD, PodTemplate
 from flyte._protos.workflow import task_definition_pb2
 from flyte._secret import SecretRequest, secrets_from_request
 from flyte._task import AsyncFunctionTaskTemplate, TaskTemplate
@@ -121,17 +124,18 @@ def get_proto_task(task: TaskTemplate, serialize_context: SerializationContext)
     # if task.parent_env is None:
     # raise ValueError(f"Task {task.name} must have a parent environment")
 
-    #
-    # This pod will be incorrect when doing fast serialize
-    #
-    container = _get_urun_container(serialize_context, task)
+    # TODO Add support for SQL, extra_config, custom
+    extra_config: typing.Dict[str, str] = {}
+    custom = {}  # type: ignore
 
-    # TODO Add support for SQL, Pod, extra_config, custom
-    pod = None
     sql = None
-    # pod = task.get_k8s_pod(serialize_context)
-    extra_config = {}  # type: ignore
-    custom = {}  # type: ignore
+    if task.pod_template and not isinstance(task.pod_template, str):
+        container = None
+        pod = _get_k8s_pod(_get_urun_container(serialize_context, task), task.pod_template)
+        extra_config[_PRIMARY_CONTAINER_NAME_FIELD] = task.pod_template.primary_container_name
+    else:
+        container = _get_urun_container(serialize_context, task)
+        pod = None
 
     # -------------- CACHE HANDLING ----------------------
     task_cache = cache_from_request(task.cache)
@@ -210,6 +214,72 @@ def _get_urun_container(
     )
 
 
+def _sanitize_resource_name(resource: tasks_pb2.Resources.ResourceEntry) -> str:
+    return tasks_pb2.Resources.ResourceName.Name(resource.name).lower().replace("_", "-")
+
+
+def _get_k8s_pod(primary_container: tasks_pb2.Container, pod_template: PodTemplate) -> Optional[tasks_pb2.K8sPod]:
+    """
+    Get the K8sPod representation of the task template.
+    :param task: The task to convert.
+    :return: The K8sPod representation of the task template.
+    """
+    from kubernetes.client import ApiClient, V1PodSpec
+    from kubernetes.client.models import V1EnvVar, V1ResourceRequirements
+
+    pod_template = copy.deepcopy(pod_template)
+    containers = cast(V1PodSpec, pod_template.pod_spec).containers
+    primary_exists = False
+
+    for container in containers:
+        if container.name == pod_template.primary_container_name:
+            primary_exists = True
+            break
+
+    if not primary_exists:
+        raise ValueError(
+            "No primary container defined in the pod spec."
+            f" You must define a primary container with the name '{pod_template.primary_container_name}'."
+        )
+    final_containers = []
+
+    for container in containers:
+        # We overwrite the primary container attributes with the values given to ContainerTask.
+        # The attributes include: image, command, args, resource, and env (env is unioned)
+
+        if container.name == pod_template.primary_container_name:
+            if container.image is None:
+                # Copy the image from primary_container only if the image is not specified in the pod spec.
+                container.image = primary_container.image
+
+            container.command = list(primary_container.command)
+            container.args = list(primary_container.args)
+
+            limits, requests = {}, {}
+            for resource in primary_container.resources.limits:
+                limits[_sanitize_resource_name(resource)] = resource.value
+            for resource in primary_container.resources.requests:
+                requests[_sanitize_resource_name(resource)] = resource.value
+
+            resource_requirements = V1ResourceRequirements(limits=limits, requests=requests)
+            if len(limits) > 0 or len(requests) > 0:
+                # Important! Only copy over resource requirements if they are non-empty.
+                container.resources = resource_requirements
+
+            if primary_container.env is not None:
+                container.env = [V1EnvVar(name=e.key, value=e.value) for e in primary_container.env] + (
+                    container.env or []
+                )
+
+        final_containers.append(container)
+
+    cast(V1PodSpec, pod_template.pod_spec).containers = final_containers
+    pod_spec = ApiClient().sanitize_for_serialization(pod_template.pod_spec)
+
+    metadata = tasks_pb2.K8sObjectMetadata(labels=pod_template.labels, annotations=pod_template.annotations)
+    return tasks_pb2.K8sPod(pod_spec=pod_spec, metadata=metadata)
+
+
 def extract_code_bundle(task_spec: task_definition_pb2.TaskSpec) -> Optional[CodeBundle]:
     """
     Extract the code bundle from the task spec.

flyte/_internal/runtime/taskrunner.py

@@ -145,7 +145,7 @@ async def convert_and_run(
             return None, convert_from_native_to_error(err)
         if task.report:
             await flyte.report.flush.aio()
-        return await convert_from_native_to_outputs(out, task.native_interface), None
+        return await convert_from_native_to_outputs(out, task.native_interface, task.name), None
 
 
 async def extract_download_run_upload(

flyte/_logging.py

@@ -6,7 +6,7 @@ from typing import Optional
 
 from ._tools import ipython_check, is_in_cluster
 
-DEFAULT_LOG_LEVEL = logging.INFO
+DEFAULT_LOG_LEVEL = logging.WARNING
 
 
 def is_rich_logging_disabled() -> bool:

flyte/_pod.py

@@ -0,0 +1,19 @@
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Dict, Optional
+
+if TYPE_CHECKING:
+    from kubernetes.client import V1PodSpec
+
+
+_PRIMARY_CONTAINER_NAME_FIELD = "primary_container_name"
+_PRIMARY_CONTAINER_DEFAULT_NAME = "primary"
+
+
+@dataclass(init=True, repr=True, eq=True, frozen=False)
+class PodTemplate(object):
+    """Custom PodTemplate specification for a Task."""
+
+    pod_spec: Optional["V1PodSpec"] = field(default_factory=lambda: V1PodSpec())
+    primary_container_name: str = _PRIMARY_CONTAINER_DEFAULT_NAME
+    labels: Optional[Dict[str, str]] = None
+    annotations: Optional[Dict[str, str]] = None

flyte/_run.py

@@ -5,6 +5,7 @@ import pathlib
 import uuid
 from typing import TYPE_CHECKING, Any, Literal, Optional, Tuple, Union, cast
 
+import flyte.errors
 from flyte.errors import InitializationError
 from flyte.models import ActionID, Checkpoints, CodeBundle, RawDataPath, SerializationContext, TaskContext
 from flyte.syncify import syncify
@@ -25,6 +26,7 @@ from ._tools import ipython_check
 
 if TYPE_CHECKING:
     from flyte.remote import Run
+    from flyte.remote._task import LazyEntity
 
     from ._code_bundle import CopyFiles
 
@@ -81,8 +83,11 @@ class _Runner:
         self._run_base_dir = run_base_dir or "/tmp/base"
 
     @requires_initialization
-    async def _run_remote(self, obj: TaskTemplate[P, R], *args: P.args, **kwargs: P.kwargs) -> Run:
+    async def _run_remote(self, obj: TaskTemplate[P, R] | LazyEntity, *args: P.args, **kwargs: P.kwargs) -> Run:
+        import grpc
+
         from flyte.remote import Run
+        from flyte.remote._task import LazyEntity
 
         from ._code_bundle import build_code_bundle, build_pkl_bundle
         from ._deploy import build_images, plan_deploy
@@ -93,37 +98,47 @@ class _Runner:
 
         cfg = get_common_config()
 
-        if obj.parent_env is None:
-            raise ValueError("Task is not attached to an environment. Please attach the task to an environment.")
+        if isinstance(obj, LazyEntity):
+            task = await obj.fetch.aio()
+            task_spec = task.pb2.spec
+            inputs = await convert_from_native_to_inputs(task.interface, *args, **kwargs)
+            version = task.pb2.task_id.version
+            code_bundle = None
+        else:
+            if obj.parent_env is None:
+                raise ValueError("Task is not attached to an environment. Please attach the task to an environment")
 
-        deploy_plan = plan_deploy(cast(Environment, obj.parent_env()))
-        image_cache = await build_images(deploy_plan)
+            deploy_plan = plan_deploy(cast(Environment, obj.parent_env()))
+            image_cache = await build_images(deploy_plan)
 
-        if self._interactive_mode:
-            code_bundle = await build_pkl_bundle(
-                obj, upload_to_controlplane=not self._dry_run, copy_bundle_to=self._copy_bundle_to
-            )
-        else:
-            if self._copy_files != "none":
-                code_bundle = await build_code_bundle(
-                    from_dir=cfg.root_dir, dryrun=self._dry_run, copy_bundle_to=self._copy_bundle_to
+            if self._interactive_mode:
+                code_bundle = await build_pkl_bundle(
+                    obj, upload_to_controlplane=not self._dry_run, copy_bundle_to=self._copy_bundle_to
                 )
             else:
-                code_bundle = None
+                if self._copy_files != "none":
+                    code_bundle = await build_code_bundle(
+                        from_dir=cfg.root_dir,
+                        dryrun=self._dry_run,
+                        copy_bundle_to=self._copy_bundle_to,
+                        copy_style=self._copy_files,
+                    )
+                else:
+                    code_bundle = None
 
-        version = self._version or (
-            code_bundle.computed_version if code_bundle and code_bundle.computed_version else None
-        )
-        if not version:
-            raise ValueError("Version is required when running a task")
-        s_ctx = SerializationContext(
-            code_bundle=code_bundle,
-            version=version,
-            image_cache=image_cache,
-            root_dir=cfg.root_dir,
-        )
-        task_spec = translate_task_to_wire(obj, s_ctx)
-        inputs = await convert_from_native_to_inputs(obj.native_interface, *args, **kwargs)
+            version = self._version or (
+                code_bundle.computed_version if code_bundle and code_bundle.computed_version else None
+            )
+            if not version:
+                raise ValueError("Version is required when running a task")
+            s_ctx = SerializationContext(
+                code_bundle=code_bundle,
+                version=version,
+                image_cache=image_cache,
+                root_dir=cfg.root_dir,
+            )
+            task_spec = translate_task_to_wire(obj, s_ctx)
+            inputs = await convert_from_native_to_inputs(obj.native_interface, *args, **kwargs)
 
         if not self._dry_run:
             if get_client() is None:
@@ -160,15 +175,35 @@ class _Runner:
             if task_spec.task_template.id.version == "":
                 task_spec.task_template.id.version = version
 
-            resp = await get_client().run_service.CreateRun(
-                run_service_pb2.CreateRunRequest(
-                    run_id=run_id,
-                    project_id=project_id,
-                    task_spec=task_spec,
-                    inputs=inputs.proto_inputs,
-                ),
-            )
-            return Run(pb2=resp.run)
+            try:
+                resp = await get_client().run_service.CreateRun(
+                    run_service_pb2.CreateRunRequest(
+                        run_id=run_id,
+                        project_id=project_id,
+                        task_spec=task_spec,
+                        inputs=inputs.proto_inputs,
+                    ),
+                )
+                return Run(pb2=resp.run)
+            except grpc.aio.AioRpcError as e:
+                if e.code() == grpc.StatusCode.UNAVAILABLE:
+                    raise flyte.errors.RuntimeSystemError(
+                        "SystemUnavailableError",
+                        "Flyte system is currently unavailable. check your configuration, or the service status.",
+                    ) from e
+                elif e.code() == grpc.StatusCode.INVALID_ARGUMENT:
+                    raise flyte.errors.RuntimeUserError("InvalidArgumentError", e.details())
+                elif e.code() == grpc.StatusCode.ALREADY_EXISTS:
+                    # TODO maybe this should be a pass and return existing run?
+                    raise flyte.errors.RuntimeUserError(
+                        "RunAlreadyExistsError",
+                        f"A run with the name '{self._name}' already exists. Please choose a different name.",
+                    )
+                else:
+                    raise flyte.errors.RuntimeSystemError(
+                        "RunCreationError",
+                        f"Failed to create run: {e.details()}",
+                    ) from e
 
         class DryRun(Run):
             def __init__(self, _task_spec, _inputs, _code_bundle):
@@ -225,7 +260,10 @@ class _Runner:
             else:
                 if self._copy_files != "none":
                     code_bundle = await build_code_bundle(
-                        from_dir=cfg.root_dir, dryrun=self._dry_run, copy_bundle_to=self._copy_bundle_to
+                        from_dir=cfg.root_dir,
+                        dryrun=self._dry_run,
+                        copy_bundle_to=self._copy_bundle_to,
+                        copy_style=self._copy_files,
                     )
                 else:
                     code_bundle = None
@@ -313,7 +351,7 @@ class _Runner:
             report=Report(name=action.name),
             mode="local",
         )
-        async with ctx.replace_task_context(tctx):
+        with ctx.replace_task_context(tctx):
             # make the local version always runs on a different thread, returns a wrapped future.
             if obj._call_as_synchronous:
                 fut = controller.submit_sync(obj, *args, **kwargs)
@@ -323,7 +361,9 @@ class _Runner:
                 return await controller.submit(obj, *args, **kwargs)
 
     @syncify
-    async def run(self, task: TaskTemplate[P, Union[R, Run]], *args: P.args, **kwargs: P.kwargs) -> Union[R, Run]:
+    async def run(
+        self, task: TaskTemplate[P, Union[R, Run]] | LazyEntity, *args: P.args, **kwargs: P.kwargs
+    ) -> Union[R, Run]:
         """
         Run an async `@env.task` or `TaskTemplate` instance. The existing async context will be used.
 
@@ -345,8 +385,13 @@ class _Runner:
         :param kwargs: Keyword arguments to pass to the Task
         :return: Run instance or the result of the task
         """
+        from flyte.remote._task import LazyEntity
+
+        if isinstance(task, LazyEntity) and self._mode != "remote":
+            raise ValueError("Remote task can only be run in remote mode.")
         if self._mode == "remote":
             return await self._run_remote(task, *args, **kwargs)
+        task = cast(TaskTemplate, task)
         if self._mode == "hybrid":
             return await self._run_hybrid(task, *args, **kwargs)
 

flyte/_task.py

@@ -23,6 +23,7 @@ from typing import (
 
 from flyteidl.core.tasks_pb2 import DataLoadingConfig
 
+from flyte._pod import PodTemplate
 from flyte.errors import RuntimeSystemError, RuntimeUserError
 
 from ._cache import Cache, CacheRequest
@@ -37,8 +38,6 @@ from ._timeout import TimeoutType
 from .models import NativeInterface, SerializationContext
 
 if TYPE_CHECKING:
-    from kubernetes.client import V1PodTemplate
-
     from ._task_environment import TaskEnvironment
 
 P = ParamSpec("P")  # capture the function's parameters
@@ -96,8 +95,7 @@ class TaskTemplate(Generic[P, R]):
     env: Optional[Dict[str, str]] = None
     secrets: Optional[SecretRequest] = None
     timeout: Optional[TimeoutType] = None
-    primary_container_name: str = "primary"
-    pod_template: Optional[Union[str, V1PodTemplate]] = None
+    pod_template: Optional[Union[str, PodTemplate]] = None
     report: bool = False
 
     parent_env: Optional[weakref.ReferenceType[TaskEnvironment]] = None
@@ -108,15 +106,6 @@ class TaskTemplate(Generic[P, R]):
     _call_as_synchronous: bool = False
 
     def __post_init__(self):
-        # If pod_template is set to a pod, verify
-        if self.pod_template is not None and not isinstance(self.pod_template, str):
-            try:
-                from kubernetes.client import V1PodTemplate  # noqa: F401
-            except ImportError as e:
-                raise ImportError(
-                    "kubernetes is not installed, please install kubernetes package to use pod_template"
-                ) from e
-
         # Auto set the image based on the image request
         if self.image == "auto":
             self.image = Image.auto()