Async/tasks remote (#2964)

auto detect flyte remote connection

Signed-off-by: Yee Hing Tong <[email protected]>

flytekit/configuration/plugin.py

@@ -18,6 +18,7 @@
 ```
 """
 
+import os
 from typing import Optional, Protocol, runtime_checkable
 
 from click import Group
@@ -59,10 +60,22 @@ def get_remote(
         config: Optional[str], project: str, domain: str, data_upload_location: Optional[str] = None
     ) -> FlyteRemote:
         """Get FlyteRemote object for CLI session."""
+
         cfg_file = get_config_file(config)
+
+        # The assumption here (if there's no config file that means we want sandbox) is too broad.
+        # todo: can improve this in the future, rather than just checking one env var, auto() with
+        #   nothing configured should probably not return sandbox but can consider
         if cfg_file is None:
-            cfg_obj = Config.for_sandbox()
-            logger.info("No config files found, creating remote with sandbox config")
+            # We really are just looking for endpoint, client_id, and client_secret. These correspond to the env vars
+            # FLYTE_PLATFORM_URL, FLYTE_CREDENTIALS_CLIENT_ID, FLYTE_CREDENTIALS_CLIENT_SECRET
+            # auto() should pick these up.
+            if "FLYTE_PLATFORM_URL" in os.environ:
+                cfg_obj = Config.auto(None)
+                logger.warning(f"Auto-created config object to pick up env vars {cfg_obj}")
+            else:
+                cfg_obj = Config.for_sandbox()
+                logger.info("No config files found, creating remote with sandbox config")
         else:  # pragma: no cover
             cfg_obj = Config.auto(config)
             logger.debug(f"Creating remote with config {cfg_obj}" + (f" with file {config}" if config else ""))

flytekit/core/base_task.py

@@ -141,6 +141,7 @@ class TaskMetadata(object):
     retries: int = 0
     timeout: Optional[Union[datetime.timedelta, int]] = None
     pod_template_name: Optional[str] = None
+    is_eager: bool = False
 
     def __post_init__(self):
         if self.timeout:
@@ -180,6 +181,7 @@ def to_taskmetadata_model(self) -> _task_model.TaskMetadata:
             cache_serializable=self.cache_serialize,
             pod_template_name=self.pod_template_name,
             cache_ignore_input_vars=self.cache_ignore_input_vars,
+            is_eager=self.is_eager,
         )
 
 

flytekit/core/python_function_task.py

@@ -28,7 +28,7 @@
 
 from flytekit.configuration import ImageConfig, SerializationSettings
 from flytekit.core import launch_plan as _annotated_launch_plan
-from flytekit.core.base_task import Task, TaskResolverMixin
+from flytekit.core.base_task import Task, TaskMetadata, TaskResolverMixin
 from flytekit.core.constants import EAGER_ROOT_ENV_NAME
 from flytekit.core.context_manager import ExecutionState, FlyteContext, FlyteContextManager
 from flytekit.core.docstring import Docstring
@@ -455,6 +455,11 @@ def __init__(
         if "execution_mode" in kwargs:
             del kwargs["execution_mode"]
 
+        if "metadata" in kwargs:
+            kwargs["metadata"].is_eager = True
+        else:
+            kwargs["metadata"] = TaskMetadata(is_eager=True)
+
         super().__init__(
             task_config,
             task_function,
@@ -516,25 +521,38 @@ async def async_execute(self, *args, **kwargs) -> Any:
                 return await self._task_function(**kwargs)
 
     def execute(self, **kwargs) -> Any:
-        from flytekit.experimental.eager_function import _internal_demo_remote
-        from flytekit.remote.remote import FlyteRemote
-
-        remote = FlyteRemote.for_sandbox(default_project="flytesnacks", default_domain="development")
-        remote = _internal_demo_remote(remote)
-
         ctx = FlyteContextManager.current_context()
         is_local_execution = cast(ExecutionState, ctx.execution_state).is_local_execution()
         builder = ctx.new_builder()
         if not is_local_execution:
             # ensure that the worker queue is in context
             if not ctx.worker_queue:
+                from flytekit.configuration.plugin import get_plugin
+
                 # This should be read from transport at real runtime if available, but if not, we should either run
                 # remote in interactive mode, or let users configure the version to use.
                 ss = ctx.serialization_settings
                 if not ss:
                     ss = SerializationSettings(
                         image_config=ImageConfig.auto_default_image(),
                     )
+
+                # In order to build the controller, we really just need a remote.
+                project = (
+                    ctx.user_space_params.execution_id.project
+                    if ctx.user_space_params and ctx.user_space_params.execution_id
+                    else "flytesnacks"
+                )
+                domain = (
+                    ctx.user_space_params.execution_id.domain
+                    if ctx.user_space_params and ctx.user_space_params.execution_id
+                    else "development"
+                )
+                raw_output = ctx.user_space_params.raw_output_prefix
+                remote = get_plugin().get_remote(
+                    config=None, project=project, domain=domain, data_upload_location=raw_output
+                )
+
                 # tag is the current execution id
                 # root tag is read from the environment variable if it exists, if not, it's the current execution id
                 if not ctx.user_space_params or not ctx.user_space_params.execution_id:
@@ -589,14 +607,3 @@ async def run_with_backend(self, **kwargs):
                 # now have to fail this eager task, because we don't want it to show up as succeeded.
                 raise FlyteNonRecoverableSystemException(base_error)
             return result
-
-
-"""
-update code comments and remove int test for now
-
-verify auth env var and start auto loading
-  - figure out how remotes can be different.
-pure watch informer pattern
-
-priority for flytekit - fix naming, depending on src
-"""

flytekit/models/task.py

@@ -179,6 +179,7 @@ def __init__(
         cache_serializable,
         pod_template_name,
         cache_ignore_input_vars,
+        is_eager: bool = False,
     ):
         """
         Information needed at runtime to determine behavior such as whether or not outputs are discoverable, timeouts,
@@ -200,6 +201,7 @@ def __init__(
             single instance over identical inputs is executed, other concurrent executions wait for the cached results.
         :param pod_template_name: The name of the existing PodTemplate resource which will be used in this task.
         :param cache_ignore_input_vars: Input variables that should not be included when calculating hash for cache.
+        :param is_eager:
         """
         self._discoverable = discoverable
         self._runtime = runtime
@@ -211,6 +213,11 @@ def __init__(
         self._cache_serializable = cache_serializable
         self._pod_template_name = pod_template_name
         self._cache_ignore_input_vars = cache_ignore_input_vars
+        self._is_eager = is_eager
+
+    @property
+    def is_eager(self):
+        return self._is_eager
 
     @property
     def discoverable(self):
@@ -310,13 +317,14 @@ def to_flyte_idl(self):
             cache_serializable=self.cache_serializable,
             pod_template_name=self.pod_template_name,
             cache_ignore_input_vars=self.cache_ignore_input_vars,
+            is_eager=self.is_eager,
         )
         if self.timeout:
             tm.timeout.FromTimedelta(self.timeout)
         return tm
 
     @classmethod
-    def from_flyte_idl(cls, pb2_object):
+    def from_flyte_idl(cls, pb2_object: _core_task.TaskMetadata):
         """
         :param flyteidl.core.task_pb2.TaskMetadata pb2_object:
         :rtype: TaskMetadata
@@ -332,6 +340,7 @@ def from_flyte_idl(cls, pb2_object):
             cache_serializable=pb2_object.cache_serializable,
             pod_template_name=pb2_object.pod_template_name,
             cache_ignore_input_vars=pb2_object.cache_ignore_input_vars,
+            is_eager=pb2_object.is_eager,
         )
 
 

pyproject.toml

@@ -20,7 +20,7 @@ dependencies = [
     "diskcache>=5.2.1",
     "docker>=4.0.0",
     "docstring-parser>=0.9.0",
-    "flyteidl>=1.13.7",
+#    "flyteidl>=1.13.7",  # todo:async bump after releasing idl
     "fsspec>=2023.3.0",
     "gcsfs>=2023.3.0",
     "googleapis-common-protos>=1.57",

tests/flytekit/unit/core/test_async.py

@@ -3,16 +3,27 @@
 from flytekit.core.worker_queue import Controller
 from flytekit.utils.asyn import loop_manager
 from flytekit.core.context_manager import FlyteContextManager
-from flytekit.configuration import Config, DataConfig, S3Config, FastSerializationSettings
+from flytekit.configuration import Config, DataConfig, S3Config, FastSerializationSettings, ImageConfig, SerializationSettings, Image
 from flytekit.core.data_persistence import FileAccessProvider
+from flytekit.tools.translator import get_serializable
+from collections import OrderedDict
+
+default_img = Image(name="default", fqn="test", tag="tag")
+serialization_settings = SerializationSettings(
+    project="project",
+    domain="domain",
+    version="version",
+    env=None,
+    image_config=ImageConfig(default_image=default_img, images=[default_img]),
+)
 
 
 @task
 def add_one(x: int) -> int:
     return x + 1
 
 
-@eager
+@eager(environment={"a": "b"})
 async def simple_eager_workflow(x: int) -> int:
     # This is the normal way of calling tasks. Call normal tasks in an effectively async way by hanging and waiting for
     # the result.
@@ -41,3 +52,9 @@ def test_easy_2():
     ):
         res = loop_manager.run_sync(simple_eager_workflow.run_with_backend, x=1)
         assert res == 2
+
+
+def test_serialization():
+    se_spec = get_serializable(OrderedDict(), serialization_settings, simple_eager_workflow)
+    assert se_spec.template.metadata.is_eager
+    assert len(se_spec.template.container.env) == 2