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

flyte/_protos/workflow/task_definition_pb2.py

@@ -13,6 +13,7 @@ _sym_db = _symbol_database.Default()
 
 from flyte._protos.common import identifier_pb2 as common_dot_identifier__pb2
 from flyte._protos.common import identity_pb2 as common_dot_identity__pb2
+from flyteidl.core import interface_pb2 as flyteidl_dot_core_dot_interface__pb2
 from flyteidl.core import tasks_pb2 as flyteidl_dot_core_dot_tasks__pb2
 from google.protobuf import timestamp_pb2 as google_dot_protobuf_dot_timestamp__pb2
 from flyte._protos.validate.validate import validate_pb2 as validate_dot_validate__pb2
@@ -20,7 +21,7 @@ from flyte._protos.workflow import common_pb2 as workflow_dot_common__pb2
 from flyte._protos.workflow import environment_pb2 as workflow_dot_environment__pb2
 
 
-DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1eworkflow/task_definition.proto\x12\x11\x63loudidl.workflow\x1a\x17\x63ommon/identifier.proto\x1a\x15\x63ommon/identity.proto\x1a\x19\x66lyteidl/core/tasks.proto\x1a\x1fgoogle/protobuf/timestamp.proto\x1a\x17validate/validate.proto\x1a\x15workflow/common.proto\x1a\x1aworkflow/environment.proto\"\x8f\x01\n\x08TaskName\x12\x1b\n\x03org\x18\x01 \x01(\tB\t\xfa\x42\x06r\x04\x10\x01\x18?R\x03org\x12#\n\x07project\x18\x02 \x01(\tB\t\xfa\x42\x06r\x04\x10\x01\x18?R\x07project\x12!\n\x06\x64omain\x18\x03 \x01(\tB\t\xfa\x42\x06r\x04\x10\x01\x18?R\x06\x64omain\x12\x1e\n\x04name\x18\x04 \x01(\tB\n\xfa\x42\x07r\x05\x10\x01\x18\xff\x01R\x04name\"\xba\x01\n\x0eTaskIdentifier\x12\x1b\n\x03org\x18\x01 \x01(\tB\t\xfa\x42\x06r\x04\x10\x01\x18?R\x03org\x12#\n\x07project\x18\x02 \x01(\tB\t\xfa\x42\x06r\x04\x10\x01\x18?R\x07project\x12!\n\x06\x64omain\x18\x03 \x01(\tB\t\xfa\x42\x06r\x04\x10\x01\x18?R\x06\x64omain\x12\x1e\n\x04name\x18\x04 \x01(\tB\n\xfa\x42\x07r\x05\x10\x01\x18\xff\x01R\x04name\x12#\n\x07version\x18\x05 \x01(\tB\t\xfa\x42\x06r\x04\x10\x01\x18?R\x07version\"\xed\x01\n\x0cTaskMetadata\x12L\n\x0b\x64\x65ployed_by\x18\x01 \x01(\x0b\x32!.cloudidl.common.EnrichedIdentityB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\ndeployedBy\x12\x1d\n\nshort_name\x18\x02 \x01(\tR\tshortName\x12\x45\n\x0b\x64\x65ployed_at\x18\x03 \x01(\x0b\x32\x1a.google.protobuf.TimestampB\x08\xfa\x42\x05\xb2\x01\x02\x08\x01R\ndeployedAt\x12)\n\x10\x65nvironment_name\x18\x04 \x01(\tR\x0f\x65nvironmentName\"\x93\x01\n\x04Task\x12\x44\n\x07task_id\x18\x01 \x01(\x0b\x32!.cloudidl.workflow.TaskIdentifierB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x06taskId\x12\x45\n\x08metadata\x18\x02 \x01(\x0b\x32\x1f.cloudidl.workflow.TaskMetadataB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x08metadata\"\x8a\x02\n\x08TaskSpec\x12J\n\rtask_template\x18\x01 \x01(\x0b\x32\x1b.flyteidl.core.TaskTemplateB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x0ctaskTemplate\x12H\n\x0e\x64\x65\x66\x61ult_inputs\x18\x02 \x03(\x0b\x32!.cloudidl.workflow.NamedParameterR\rdefaultInputs\x12&\n\nshort_name\x18\x03 \x01(\tB\x07\xfa\x42\x04r\x02\x18?R\tshortName\x12@\n\x0b\x65nvironment\x18\x04 \x01(\x0b\x32\x1e.cloudidl.workflow.EnvironmentR\x0b\x65nvironment\"\xd5\x01\n\x0bTaskDetails\x12\x44\n\x07task_id\x18\x01 \x01(\x0b\x32!.cloudidl.workflow.TaskIdentifierB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x06taskId\x12\x45\n\x08metadata\x18\x02 \x01(\x0b\x32\x1f.cloudidl.workflow.TaskMetadataB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x08metadata\x12\x39\n\x04spec\x18\x03 \x01(\x0b\x32\x1b.cloudidl.workflow.TaskSpecB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x04specB\xc0\x01\n\x15\x63om.cloudidl.workflowB\x13TaskDefinitionProtoH\x02P\x01Z+github.com/unionai/cloud/gen/pb-go/workflow\xa2\x02\x03\x43WX\xaa\x02\x11\x43loudidl.Workflow\xca\x02\x11\x43loudidl\\Workflow\xe2\x02\x1d\x43loudidl\\Workflow\\GPBMetadata\xea\x02\x12\x43loudidl::Workflowb\x06proto3')
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1eworkflow/task_definition.proto\x12\x11\x63loudidl.workflow\x1a\x17\x63ommon/identifier.proto\x1a\x15\x63ommon/identity.proto\x1a\x1d\x66lyteidl/core/interface.proto\x1a\x19\x66lyteidl/core/tasks.proto\x1a\x1fgoogle/protobuf/timestamp.proto\x1a\x17validate/validate.proto\x1a\x15workflow/common.proto\x1a\x1aworkflow/environment.proto\"\x8f\x01\n\x08TaskName\x12\x1b\n\x03org\x18\x01 \x01(\tB\t\xfa\x42\x06r\x04\x10\x01\x18?R\x03org\x12#\n\x07project\x18\x02 \x01(\tB\t\xfa\x42\x06r\x04\x10\x01\x18?R\x07project\x12!\n\x06\x64omain\x18\x03 \x01(\tB\t\xfa\x42\x06r\x04\x10\x01\x18?R\x06\x64omain\x12\x1e\n\x04name\x18\x04 \x01(\tB\n\xfa\x42\x07r\x05\x10\x01\x18\xff\x01R\x04name\"\xba\x01\n\x0eTaskIdentifier\x12\x1b\n\x03org\x18\x01 \x01(\tB\t\xfa\x42\x06r\x04\x10\x01\x18?R\x03org\x12#\n\x07project\x18\x02 \x01(\tB\t\xfa\x42\x06r\x04\x10\x01\x18?R\x07project\x12!\n\x06\x64omain\x18\x03 \x01(\tB\t\xfa\x42\x06r\x04\x10\x01\x18?R\x06\x64omain\x12\x1e\n\x04name\x18\x04 \x01(\tB\n\xfa\x42\x07r\x05\x10\x01\x18\xff\x01R\x04name\x12#\n\x07version\x18\x05 \x01(\tB\t\xfa\x42\x06r\x04\x10\x01\x18?R\x07version\"\xed\x01\n\x0cTaskMetadata\x12L\n\x0b\x64\x65ployed_by\x18\x01 \x01(\x0b\x32!.cloudidl.common.EnrichedIdentityB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\ndeployedBy\x12\x1d\n\nshort_name\x18\x02 \x01(\tR\tshortName\x12\x45\n\x0b\x64\x65ployed_at\x18\x03 \x01(\x0b\x32\x1a.google.protobuf.TimestampB\x08\xfa\x42\x05\xb2\x01\x02\x08\x01R\ndeployedAt\x12)\n\x10\x65nvironment_name\x18\x04 \x01(\tR\x0f\x65nvironmentName\"\x93\x01\n\x04Task\x12\x44\n\x07task_id\x18\x01 \x01(\x0b\x32!.cloudidl.workflow.TaskIdentifierB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x06taskId\x12\x45\n\x08metadata\x18\x02 \x01(\x0b\x32\x1f.cloudidl.workflow.TaskMetadataB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x08metadata\"\x8a\x02\n\x08TaskSpec\x12J\n\rtask_template\x18\x01 \x01(\x0b\x32\x1b.flyteidl.core.TaskTemplateB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x0ctaskTemplate\x12H\n\x0e\x64\x65\x66\x61ult_inputs\x18\x02 \x03(\x0b\x32!.cloudidl.workflow.NamedParameterR\rdefaultInputs\x12&\n\nshort_name\x18\x03 \x01(\tB\x07\xfa\x42\x04r\x02\x18?R\tshortName\x12@\n\x0b\x65nvironment\x18\x04 \x01(\x0b\x32\x1e.cloudidl.workflow.EnvironmentR\x0b\x65nvironment\"H\n\tTraceSpec\x12;\n\tinterface\x18\x01 \x01(\x0b\x32\x1d.flyteidl.core.TypedInterfaceR\tinterface\"\xd5\x01\n\x0bTaskDetails\x12\x44\n\x07task_id\x18\x01 \x01(\x0b\x32!.cloudidl.workflow.TaskIdentifierB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x06taskId\x12\x45\n\x08metadata\x18\x02 \x01(\x0b\x32\x1f.cloudidl.workflow.TaskMetadataB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x08metadata\x12\x39\n\x04spec\x18\x03 \x01(\x0b\x32\x1b.cloudidl.workflow.TaskSpecB\x08\xfa\x42\x05\x8a\x01\x02\x10\x01R\x04specB\xc0\x01\n\x15\x63om.cloudidl.workflowB\x13TaskDefinitionProtoH\x02P\x01Z+github.com/unionai/cloud/gen/pb-go/workflow\xa2\x02\x03\x43WX\xaa\x02\x11\x43loudidl.Workflow\xca\x02\x11\x43loudidl\\Workflow\xe2\x02\x1d\x43loudidl\\Workflow\\GPBMetadata\xea\x02\x12\x43loudidl::Workflowb\x06proto3')
 
 _globals = globals()
 _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
@@ -64,16 +65,18 @@ if _descriptor._USE_C_DESCRIPTORS == False:
   _TASKDETAILS.fields_by_name['metadata']._serialized_options = b'\372B\005\212\001\002\020\001'
   _TASKDETAILS.fields_by_name['spec']._options = None
   _TASKDETAILS.fields_by_name['spec']._serialized_options = b'\372B\005\212\001\002\020\001'
-  _globals['_TASKNAME']._serialized_start=238
-  _globals['_TASKNAME']._serialized_end=381
-  _globals['_TASKIDENTIFIER']._serialized_start=384
-  _globals['_TASKIDENTIFIER']._serialized_end=570
-  _globals['_TASKMETADATA']._serialized_start=573
-  _globals['_TASKMETADATA']._serialized_end=810
-  _globals['_TASK']._serialized_start=813
-  _globals['_TASK']._serialized_end=960
-  _globals['_TASKSPEC']._serialized_start=963
-  _globals['_TASKSPEC']._serialized_end=1229
-  _globals['_TASKDETAILS']._serialized_start=1232
-  _globals['_TASKDETAILS']._serialized_end=1445
+  _globals['_TASKNAME']._serialized_start=269
+  _globals['_TASKNAME']._serialized_end=412
+  _globals['_TASKIDENTIFIER']._serialized_start=415
+  _globals['_TASKIDENTIFIER']._serialized_end=601
+  _globals['_TASKMETADATA']._serialized_start=604
+  _globals['_TASKMETADATA']._serialized_end=841
+  _globals['_TASK']._serialized_start=844
+  _globals['_TASK']._serialized_end=991
+  _globals['_TASKSPEC']._serialized_start=994
+  _globals['_TASKSPEC']._serialized_end=1260
+  _globals['_TRACESPEC']._serialized_start=1262
+  _globals['_TRACESPEC']._serialized_end=1334
+  _globals['_TASKDETAILS']._serialized_start=1337
+  _globals['_TASKDETAILS']._serialized_end=1550
 # @@protoc_insertion_point(module_scope)

flyte/_protos/workflow/task_definition_pb2.pyi

@@ -1,5 +1,6 @@
 from flyte._protos.common import identifier_pb2 as _identifier_pb2
 from flyte._protos.common import identity_pb2 as _identity_pb2
+from flyteidl.core import interface_pb2 as _interface_pb2
 from flyteidl.core import tasks_pb2 as _tasks_pb2
 from google.protobuf import timestamp_pb2 as _timestamp_pb2
 from flyte._protos.validate.validate import validate_pb2 as _validate_pb2
@@ -70,6 +71,12 @@ class TaskSpec(_message.Message):
     environment: _environment_pb2.Environment
     def __init__(self, task_template: _Optional[_Union[_tasks_pb2.TaskTemplate, _Mapping]] = ..., default_inputs: _Optional[_Iterable[_Union[_common_pb2.NamedParameter, _Mapping]]] = ..., short_name: _Optional[str] = ..., environment: _Optional[_Union[_environment_pb2.Environment, _Mapping]] = ...) -> None: ...
 
+class TraceSpec(_message.Message):
+    __slots__ = ["interface"]
+    INTERFACE_FIELD_NUMBER: _ClassVar[int]
+    interface: _interface_pb2.TypedInterface
+    def __init__(self, interface: _Optional[_Union[_interface_pb2.TypedInterface, _Mapping]] = ...) -> None: ...
+
 class TaskDetails(_message.Message):
     __slots__ = ["task_id", "metadata", "spec"]
     TASK_ID_FIELD_NUMBER: _ClassVar[int]

flyte/_task.py

@@ -85,7 +85,7 @@ class TaskTemplate(Generic[P, R]):
 
     name: str
     interface: NativeInterface
-    friendly_name: str = ""
+    short_name: str = ""
     task_type: str = "python"
     task_type_version: int = 0
     image: Union[str, Image, Literal["auto"]] = "auto"
@@ -129,9 +129,9 @@ class TaskTemplate(Generic[P, R]):
         if isinstance(self.retries, int):
             self.retries = RetryStrategy(count=self.retries)
 
-        if self.friendly_name == "":
-            # If friendly_name is not set, use the name of the task
-            self.friendly_name = self.name
+        if self.short_name == "":
+            # If short_name is not set, use the name of the task
+            self.short_name = self.name
 
     def __getstate__(self):
         """
@@ -314,7 +314,7 @@ class TaskTemplate(Generic[P, R]):
     def override(
         self,
         *,
-        friendly_name: Optional[str] = None,
+        short_name: Optional[str] = None,
         resources: Optional[Resources] = None,
         cache: Optional[CacheRequest] = None,
         retries: Union[int, RetryStrategy] = 0,
@@ -375,7 +375,7 @@ class TaskTemplate(Generic[P, R]):
 
         return replace(
             self,
-            friendly_name=friendly_name or self.friendly_name,
+            short_name=short_name or self.short_name,
             resources=resources,
             cache=cache,
             retries=retries,

flyte/_task_environment.py

@@ -133,7 +133,7 @@ class TaskEnvironment(Environment):
         self,
         _func=None,
         *,
-        name: Optional[str] = None,
+        short_name: Optional[str] = None,
         cache: CacheRequest | None = None,
         retries: Union[int, RetryStrategy] = 0,
         timeout: Union[timedelta, int] = 0,
@@ -147,7 +147,7 @@ class TaskEnvironment(Environment):
 
         :param _func: Optional The function to decorate. If not provided, the decorator will return a callable that
         accepts a function to be decorated.
-        :param name: Optional A friendly name for the task (defaults to the function name)
+        :param short_name: Optional A friendly name for the task (defaults to the function name)
         :param cache: Optional The cache policy for the task, defaults to auto, which will cache the results of the
         task.
         :param retries: Optional The number of retries for the task, defaults to 0, which means no retries.
@@ -166,7 +166,7 @@ class TaskEnvironment(Environment):
                 raise ValueError("Cannot set pod_template when environment is reusable.")
 
         def decorator(func: FunctionTypes) -> AsyncFunctionTaskTemplate[P, R]:
-            friendly_name = name or func.__name__
+            short = short_name or func.__name__
             task_name = self.name + "." + func.__name__
 
             if not inspect.iscoroutinefunction(func) and self.reusable is not None:
@@ -207,7 +207,7 @@ class TaskEnvironment(Environment):
                 parent_env=weakref.ref(self),
                 interface=NativeInterface.from_callable(func),
                 report=report,
-                friendly_name=friendly_name,
+                short_name=short,
                 plugin_config=self.plugin_config,
                 max_inline_io_bytes=max_inline_io_bytes,
             )

flyte/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '2.0.0b13'
-__version_tuple__ = version_tuple = (2, 0, 0, 'b13')
+__version__ = version = '2.0.0b15'
+__version_tuple__ = version_tuple = (2, 0, 0, 'b15')
 
-__commit_id__ = commit_id = 'g07f30e36d'
+__commit_id__ = commit_id = 'gaa61639b5'

flyte/cli/_build.py

@@ -3,8 +3,7 @@ from pathlib import Path
 from types import ModuleType
 from typing import Any, Dict, List, cast
 
-import click
-from click import Context
+import rich_click as click
 
 import flyte
 
@@ -44,7 +43,7 @@ class BuildEnvCommand(click.Command):
         self.build_args = build_args
         super().__init__(*args, **kwargs)
 
-    def invoke(self, ctx: Context):
+    def invoke(self, ctx: click.Context):
         from rich.console import Console
 
         console = Console()

flyte/cli/_run.py

@@ -8,8 +8,7 @@ from pathlib import Path
 from types import ModuleType
 from typing import Any, Dict, List, cast
 
-import click
-from click import Context, Parameter
+import rich_click as click
 from rich.console import Console
 from typing_extensions import get_args
 
@@ -24,7 +23,7 @@ RUN_REMOTE_CMD = "deployed-task"
 
 
 @lru_cache()
-def _initialize_config(ctx: Context, project: str, domain: str):
+def _initialize_config(ctx: click.Context, project: str, domain: str):
     obj: CLIConfig | None = ctx.obj
     if obj is None:
         import flyte.config
@@ -37,7 +36,7 @@ def _initialize_config(ctx: Context, project: str, domain: str):
 
 @lru_cache()
 def _list_tasks(
-    ctx: Context,
+    ctx: click.Context,
     project: str,
     domain: str,
     by_task_name: str | None = None,
@@ -121,7 +120,7 @@ class RunTaskCommand(click.Command):
         kwargs.pop("name", None)
         super().__init__(obj_name, *args, **kwargs)
 
-    def invoke(self, ctx: Context):
+    def invoke(self, ctx: click.Context):
         obj: CLIConfig = _initialize_config(ctx, self.run_args.project, self.run_args.domain)
 
         async def _run():
@@ -152,7 +151,7 @@ class RunTaskCommand(click.Command):
 
         asyncio.run(_run())
 
-    def get_params(self, ctx: Context) -> List[Parameter]:
+    def get_params(self, ctx: click.Context) -> List[click.Parameter]:
         # Note this function may be called multiple times by click.
         task = self.obj
         from .._internal.runtime.types_serde import transform_native_to_typed_interface
@@ -162,7 +161,7 @@ class RunTaskCommand(click.Command):
             return super().get_params(ctx)
         inputs_interface = task.native_interface.inputs
 
-        params: List[Parameter] = []
+        params: List[click.Parameter] = []
         for name, var in interface.inputs.variables.items():
             default_val = None
             if inputs_interface[name][1] is not inspect._empty:
@@ -239,7 +238,7 @@ class RunReferenceTaskCommand(click.Command):
 
         asyncio.run(_run())
 
-    def get_params(self, ctx: Context) -> List[Parameter]:
+    def get_params(self, ctx: click.Context) -> List[click.Parameter]:
         # Note this function may be called multiple times by click.
         import flyte.remote
         from flyte._internal.runtime.types_serde import transform_native_to_typed_interface
@@ -254,7 +253,7 @@ class RunReferenceTaskCommand(click.Command):
             return super().get_params(ctx)
         inputs_interface = task_details.interface.inputs
 
-        params: List[Parameter] = []
+        params: List[click.Parameter] = []
         for name, var in interface.inputs.variables.items():
             default_val = None
             if inputs_interface[name][1] is not inspect._empty:
@@ -322,7 +321,6 @@ class ReferenceTaskGroup(common.GroupBase):
 
     def get_command(self, ctx, name):
         env, task, version = self._parse_task_name(name)
-
         match env, task, version:
             case env, None, None:
                 if self._env_is_task(ctx, env):
@@ -383,10 +381,11 @@ class TaskFiles(common.FileGroup):
         super().__init__(*args, directory=directory, **kwargs)
 
     def list_commands(self, ctx):
-        return [
+        v = [
             RUN_REMOTE_CMD,
-            *self.files,
+            *super().list_commands(ctx),
         ]
+        return v
 
     def get_command(self, ctx, cmd_name):
         run_args = RunArguments.from_dict(ctx.params)

flyte/models.py

@@ -166,6 +166,7 @@ class TaskContext:
     action: ActionID
     version: str
     raw_data_path: RawDataPath
+    input_path: str | None = None
     output_path: str
     run_base_dir: str
     report: Report
@@ -175,6 +176,7 @@ class TaskContext:
     compiled_image_cache: ImageCache | None = None
     data: Dict[str, Any] = field(default_factory=dict)
     mode: Literal["local", "remote", "hybrid"] = "remote"
+    interactive_mode: bool = False
 
     def replace(self, **kwargs) -> TaskContext:
         if "data" in kwargs:

flyte/remote/_action.py

@@ -627,8 +627,11 @@ class ActionDetails(ToJSONMixin):
             )
         )
         native_iface = None
-        if self.pb2.resolved_task_spec:
-            iface = self.pb2.resolved_task_spec.task_template.interface
+        if self.pb2.HasField('task'):
+            iface = self.pb2.task.task_template.interface
+            native_iface = types.guess_interface(iface)
+        elif self.pb2.HasField('trace'):
+            iface = self.pb2.trace.interface
             native_iface = types.guess_interface(iface)
 
         if resp.inputs:

flyte/remote/_client/auth/_authenticators/device_code.py

@@ -81,7 +81,7 @@ class DeviceCodeAuthenticator(Authenticator):
                     for_endpoint=self._endpoint,
                 )
             except (AuthenticationError, AuthenticationPending):
-                logger.warning("Failed to refresh token. Kicking off a full authorization flow.")
+                logger.warning("Logging in...")
 
         """Fall back to device flow"""
         resp = await token_client.get_device_code(

flyte/remote/_client/auth/_authenticators/pkce.py

@@ -123,7 +123,7 @@ class PKCEAuthenticator(Authenticator):
             try:
                 return await self._auth_client.refresh_access_token(self._creds)
             except AccessTokenNotFoundError:
-                logger.warning("Failed to refresh token. Kicking off a full authorization flow.")
+                logger.warning("Logging in...")
 
         return await self._auth_client.get_creds_from_remote()
 

flyte/remote/_task.py

@@ -283,7 +283,7 @@ class TaskDetails(ToJSONMixin):
     def override(
         self,
         *,
-        friendly_name: Optional[str] = None,
+        short_name: Optional[str] = None,
         resources: Optional[flyte.Resources] = None,
         retries: Union[int, flyte.RetryStrategy] = 0,
         timeout: Optional[flyte.TimeoutType] = None,
@@ -297,8 +297,8 @@ class TaskDetails(ToJSONMixin):
                 f"Check the parameters for override method."
             )
         template = self.pb2.spec.task_template
-        if friendly_name:
-            self.pb2.metadata.short_name = friendly_name
+        if short_name:
+            self.pb2.metadata.short_name = short_name
         if secrets:
             template.security_context.CopyFrom(get_security_context(secrets))
         if template.HasField("container"):
@@ -318,7 +318,7 @@ class TaskDetails(ToJSONMixin):
         """
         Rich representation of the task.
         """
-        yield "friendly_name", self.pb2.spec.short_name
+        yield "short_name", self.pb2.spec.short_name
         yield "environment", self.pb2.spec.environment
         yield "default_inputs_keys", self.default_input_args
         yield "required_args", self.required_args

flyte-2.0.0b15.data/scripts/debug.py

@@ -0,0 +1,38 @@
+import click
+
+
+@click.group()
+def _debug():
+    """Debug commands for Flyte."""
+
+
+@_debug.command("resume")
+@click.option("--pid", "-m", type=int, required=True, help="PID of the vscode server.")
+def resume(pid):
+    """
+    Resume a Flyte task for debugging purposes.
+
+    Args:
+        pid (int): PID of the vscode server.
+    """
+    import os
+    import signal
+
+    print("Terminating server and resuming task.")
+    answer = (
+        input(
+            "This operation will kill the server. All unsaved data will be lost,"
+            " and you will no longer be able to connect to it. Do you really want to terminate? (Y/N): "
+        )
+        .strip()
+        .upper()
+    )
+    if answer == "Y":
+        os.kill(pid, signal.SIGTERM)
+        print("The server has been terminated and the task has been resumed.")
+    else:
+        print("Operation canceled.")
+
+
+if __name__ == "__main__":
+    _debug()

{flyte-2.0.0b13.data → flyte-2.0.0b15.data}/scripts/runtime.py

@@ -26,6 +26,7 @@ DOMAIN_NAME = "FLYTE_INTERNAL_TASK_DOMAIN"
 ORG_NAME = "_U_ORG_NAME"
 ENDPOINT_OVERRIDE = "_U_EP_OVERRIDE"
 RUN_OUTPUT_BASE_DIR = "_U_RUN_BASE"
+FLYTE_ENABLE_VSCODE_KEY = "_F_E_VS"
 
 # TODO: Remove this after proper auth is implemented
 _UNION_EAGER_API_KEY_ENV_VAR = "_UNION_EAGER_API_KEY"
@@ -49,6 +50,8 @@ def _pass_through():
 @click.option("--project", envvar=PROJECT_NAME, required=False)
 @click.option("--domain", envvar=DOMAIN_NAME, required=False)
 @click.option("--org", envvar=ORG_NAME, required=False)
+@click.option("--debug", envvar=FLYTE_ENABLE_VSCODE_KEY, type=click.BOOL, required=False)
+@click.option("--interactive-mode", type=click.BOOL, required=False)
 @click.option("--image-cache", required=False)
 @click.option("--tgz", required=False)
 @click.option("--pkl", required=False)
@@ -59,12 +62,16 @@ def _pass_through():
     type=click.UNPROCESSED,
     nargs=-1,
 )
+@click.pass_context
 def main(
+    ctx: click.Context,
     run_name: str,
     name: str,
     project: str,
     domain: str,
     org: str,
+    debug: bool,
+    interactive_mode: bool,
     image_cache: str,
     version: str,
     inputs: str,
@@ -109,6 +116,11 @@ def main(
     if name.startswith("{{"):
         name = os.getenv("ACTION_NAME", "")
 
+    if debug and name == "a0":
+        from flyte._debug.vscode import _start_vscode_server
+
+        asyncio.run(_start_vscode_server(ctx))
+
     # Figure out how to connect
     # This detection of api key is a hack for now.
     controller_kwargs: dict[str, Any] = {"insecure": False}
@@ -143,6 +155,7 @@ def main(
         version=version,
         controller=controller,
         image_cache=ic,
+        interactive_mode=interactive_mode or debug,
     )
     # Create a coroutine to watch for errors
     controller_failure = controller.watch_for_errors()

{flyte-2.0.0b13.dist-info → flyte-2.0.0b15.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flyte
-Version: 2.0.0b13
+Version: 2.0.0b15
 Summary: Add your description here
 Author-email: Ketan Umare <kumare3@users.noreply.github.com>
 Requires-Python: >=3.10
@@ -16,7 +16,7 @@ Requires-Dist: obstore>=0.7.3
 Requires-Dist: protobuf>=6.30.1
 Requires-Dist: pydantic>=2.10.6
 Requires-Dist: pyyaml>=6.0.2
-Requires-Dist: rich-click>=1.8.9
+Requires-Dist: rich-click==1.8.9
 Requires-Dist: httpx<1.0.0,>=0.28.1
 Requires-Dist: keyring>=25.6.0
 Requires-Dist: msgpack>=1.1.0