diff --git a/dev-requirements.in b/dev-requirements.in index c3a8abe025..0fdd69b52a 100644 --- a/dev-requirements.in +++ b/dev-requirements.in @@ -45,3 +45,5 @@ pandas scikit-learn types-requests prometheus-client + +orjson diff --git a/dev-requirements.txt b/dev-requirements.txt index 0634217be7..dc929286e7 100644 --- a/dev-requirements.txt +++ b/dev-requirements.txt @@ -14,7 +14,7 @@ adlfs==2023.9.0 # via flytekit aiobotocore==2.5.4 # via s3fs -aiohttp==3.8.6 +aiohttp==3.9.2 # via # adlfs # aiobotocore @@ -69,9 +69,7 @@ cfgv==3.4.0 chardet==5.2.0 # via binaryornot charset-normalizer==3.3.2 - # via - # aiohttp - # requests + # via requests click==8.1.7 # via # cookiecutter @@ -316,6 +314,8 @@ oauthlib==3.2.2 # requests-oauthlib opt-einsum==3.3.0 # via tensorflow +orjson==3.9.12 + # via -r dev-requirements.in packaging==23.2 # via # docker @@ -329,7 +329,7 @@ parso==0.8.3 # via jedi pexpect==4.8.0 # via ipython -pillow==10.1.0 +pillow==10.2.0 # via -r dev-requirements.in platformdirs==3.11.0 # via virtualenv diff --git a/doc-requirements.txt b/doc-requirements.txt index 52cc2a5d29..5c5c660b38 100644 --- a/doc-requirements.txt +++ b/doc-requirements.txt @@ -10,7 +10,7 @@ adlfs==2023.9.0 # via flytekit aiobotocore==2.5.4 # via s3fs -aiohttp==3.9.1 +aiohttp==3.9.2 # via # adlfs # aiobotocore diff --git a/flytekit/clients/friendly.py b/flytekit/clients/friendly.py index 76a80708ba..3a516fc1d2 100644 --- a/flytekit/clients/friendly.py +++ b/flytekit/clients/friendly.py @@ -12,7 +12,6 @@ from flyteidl.admin import task_pb2 as _task_pb2 from flyteidl.admin import workflow_attributes_pb2 as _workflow_attributes_pb2 from flyteidl.admin import workflow_pb2 as _workflow_pb2 -from flyteidl.artifact import artifacts_pb2 from flyteidl.service import dataproxy_pb2 as _data_proxy_pb2 from google.protobuf.duration_pb2 import Duration @@ -1037,9 +1036,3 @@ def get_data(self, flyte_uri: str) -> _data_proxy_pb2.GetDataResponse: resp = self._dataproxy_stub.GetData(req, metadata=self._metadata) return resp - - def create_artifact(self, request: artifacts_pb2.CreateArtifactRequest) -> artifacts_pb2.CreateArtifactResponse: - return self._artifact_stub.CreateArtifact(request) - - def get_artifact(self, request: artifacts_pb2.GetArtifactRequest) -> artifacts_pb2.GetArtifactResponse: - return self._artifact_stub.GetArtifact(request) diff --git a/flytekit/clients/raw.py b/flytekit/clients/raw.py index 5ef5c930bd..681a1b0071 100644 --- a/flytekit/clients/raw.py +++ b/flytekit/clients/raw.py @@ -5,7 +5,6 @@ import grpc from flyteidl.admin.project_pb2 import ProjectListRequest from flyteidl.admin.signal_pb2 import SignalList, SignalListRequest, SignalSetRequest, SignalSetResponse -from flyteidl.artifact import artifacts_pb2_grpc as artifact_service from flyteidl.service import admin_pb2_grpc as _admin_service from flyteidl.service import dataproxy_pb2 as _dataproxy_pb2 from flyteidl.service import dataproxy_pb2_grpc as dataproxy_service @@ -53,7 +52,6 @@ def __init__(self, cfg: PlatformConfig, **kwargs): self._stub = _admin_service.AdminServiceStub(self._channel) self._signal = signal_service.SignalServiceStub(self._channel) self._dataproxy_stub = dataproxy_service.DataProxyServiceStub(self._channel) - self._artifact_stub = artifact_service.ArtifactRegistryStub(self._channel) logger.info( f"Flyte Client configured -> {self._cfg.endpoint} in {'insecure' if self._cfg.insecure else 'secure'} mode." diff --git a/flytekit/clis/sdk_in_container/register.py b/flytekit/clis/sdk_in_container/register.py index 4f8dddbc18..45e41efe47 100644 --- a/flytekit/clis/sdk_in_container/register.py +++ b/flytekit/clis/sdk_in_container/register.py @@ -117,6 +117,14 @@ callback=key_value_callback, help="Environment variables to set in the container, of the format `ENV_NAME=ENV_VALUE`", ) +@click.option( + "--skip-errors", + "--skip-error", + default=False, + is_flag=True, + help="Skip errors during registration. This is useful when registering multiple packages and you want to skip " + "errors for some packages.", +) @click.argument("package-or-module", type=click.Path(exists=True, readable=True, resolve_path=True), nargs=-1) @click.pass_context def register( @@ -135,6 +143,7 @@ def register( dry_run: bool, activate_launchplans: bool, env: typing.Optional[typing.Dict[str, str]], + skip_errors: bool, ): """ see help @@ -187,6 +196,7 @@ def register( env=env, dry_run=dry_run, activate_launchplans=activate_launchplans, + skip_errors=skip_errors, ) except Exception as e: raise e diff --git a/flytekit/clis/sdk_in_container/run.py b/flytekit/clis/sdk_in_container/run.py index d5a96e069a..e8c7e6ddf6 100644 --- a/flytekit/clis/sdk_in_container/run.py +++ b/flytekit/clis/sdk_in_container/run.py @@ -513,6 +513,9 @@ def _run(*args, **kwargs): if not run_level_params.is_remote: with FlyteContextManager.with_context(_update_flyte_context(run_level_params)): + if run_level_params.envvars: + for env_var, value in run_level_params.envvars.items(): + os.environ[env_var] = value output = entity(**inputs) if inspect.iscoroutine(output): # TODO: make eager mode workflows run with local-mode @@ -792,6 +795,8 @@ def get_command(self, ctx, exe_entity): is_workflow = False if self._entities: is_workflow = exe_entity in self._entities.workflows + if not os.path.exists(self._filename): + raise ValueError(f"File {self._filename} does not exist") rel_path = os.path.relpath(self._filename) if rel_path.startswith(".."): raise ValueError( diff --git a/flytekit/clis/sdk_in_container/serve.py b/flytekit/clis/sdk_in_container/serve.py index 6a262c91a4..87f008b084 100644 --- a/flytekit/clis/sdk_in_container/serve.py +++ b/flytekit/clis/sdk_in_container/serve.py @@ -1,7 +1,10 @@ from concurrent import futures import rich_click as click -from flyteidl.service.agent_pb2_grpc import add_AsyncAgentServiceServicer_to_server +from flyteidl.service.agent_pb2_grpc import ( + add_AgentMetadataServiceServicer_to_server, + add_AsyncAgentServiceServicer_to_server, +) from grpc import aio @@ -49,7 +52,7 @@ def agent(_: click.Context, port, worker, timeout): async def _start_grpc_server(port: int, worker: int, timeout: int): click.secho("Starting up the server to expose the prometheus metrics...", fg="blue") - from flytekit.extend.backend.agent_service import AsyncAgentService + from flytekit.extend.backend.agent_service import AgentMetadataService, AsyncAgentService try: from prometheus_client import start_http_server @@ -61,6 +64,7 @@ async def _start_grpc_server(port: int, worker: int, timeout: int): server = aio.server(futures.ThreadPoolExecutor(max_workers=worker)) add_AsyncAgentServiceServicer_to_server(AsyncAgentService(), server) + add_AgentMetadataServiceServicer_to_server(AgentMetadataService(), server) server.add_insecure_port(f"[::]:{port}") await server.start() diff --git a/flytekit/configuration/__init__.py b/flytekit/configuration/__init__.py index aa9c2bf129..853e08a166 100644 --- a/flytekit/configuration/__init__.py +++ b/flytekit/configuration/__init__.py @@ -611,6 +611,22 @@ def auto(cls, config_file: typing.Union[str, ConfigFile] = None) -> DataConfig: ) +@dataclass(init=True, repr=True, eq=True, frozen=True) +class LocalConfig(object): + """ + Any configuration specific to local runs. + """ + + cache_enabled: bool = True + + @classmethod + def auto(cls, config_file: typing.Union[str, ConfigFile] = None) -> LocalConfig: + config_file = get_config_file(config_file) + kwargs = {} + kwargs = set_if_exists(kwargs, "cache_enabled", _internal.Local.CACHE_ENABLED.read(config_file)) + return LocalConfig(**kwargs) + + @dataclass(init=True, repr=True, eq=True, frozen=True) class Config(object): """ diff --git a/flytekit/configuration/default_images.py b/flytekit/configuration/default_images.py index ea9d162a8d..380f428154 100644 --- a/flytekit/configuration/default_images.py +++ b/flytekit/configuration/default_images.py @@ -1,6 +1,7 @@ import enum import sys import typing +from contextlib import suppress class PythonVersion(enum.Enum): @@ -26,6 +27,13 @@ class DefaultImages(object): @classmethod def default_image(cls) -> str: + from flytekit.configuration.plugin import get_plugin + + with suppress(AttributeError): + default_image = get_plugin().get_default_image() + if default_image is not None: + return default_image + return cls.find_image_for() @classmethod diff --git a/flytekit/configuration/internal.py b/flytekit/configuration/internal.py index e559650e29..66e21b25d7 100644 --- a/flytekit/configuration/internal.py +++ b/flytekit/configuration/internal.py @@ -66,6 +66,11 @@ class AZURE(object): CLIENT_SECRET = ConfigEntry(LegacyConfigEntry(SECTION, "client_secret")) +class Local(object): + SECTION = "local" + CACHE_ENABLED = ConfigEntry(LegacyConfigEntry(SECTION, "cache_enabled", bool)) + + class Credentials(object): SECTION = "credentials" COMMAND = ConfigEntry(LegacyConfigEntry(SECTION, "command", list), YamlConfigEntry("admin.command", list)) diff --git a/flytekit/configuration/plugin.py b/flytekit/configuration/plugin.py index 051d421d04..e29b57f727 100644 --- a/flytekit/configuration/plugin.py +++ b/flytekit/configuration/plugin.py @@ -43,6 +43,10 @@ def configure_pyflyte_cli(main: Group) -> Group: def secret_requires_group() -> bool: """Return True if secrets require group entry.""" + @staticmethod + def get_default_image() -> Optional[str]: + """Get default image. Return None to use the images from flytekit.configuration.DefaultImages""" + class FlytekitPlugin: @staticmethod @@ -71,6 +75,11 @@ def secret_requires_group() -> bool: """Return True if secrets require group entry.""" return True + @staticmethod + def get_default_image() -> Optional[str]: + """Get default image. Return None to use the images from flytekit.configuration.DefaultImages""" + return None + def _get_plugin_from_entrypoint(): """Get plugin from entrypoint.""" diff --git a/flytekit/core/array_node_map_task.py b/flytekit/core/array_node_map_task.py index e4fef9ed10..5d10cadd45 100644 --- a/flytekit/core/array_node_map_task.py +++ b/flytekit/core/array_node_map_task.py @@ -77,8 +77,9 @@ def __init__( f = actual_task.lhs else: _, mod, f, _ = tracker.extract_task_module(cast(PythonFunctionTask, actual_task).task_function) + sorted_bounded_inputs = ",".join(sorted(self._bound_inputs)) h = hashlib.md5( - f"{collection_interface.__str__()}{concurrency}{min_successes}{min_success_ratio}".encode("utf-8") + f"{sorted_bounded_inputs}{concurrency}{min_successes}{min_success_ratio}".encode("utf-8") ).hexdigest() self._name = f"{mod}.map_{f}_{h}-arraynode" @@ -387,7 +388,7 @@ def load_task(self, loader_args: List[str], max_concurrency: int = 0) -> ArrayNo def loader_args(self, settings: SerializationSettings, t: ArrayNodeMapTask) -> List[str]: # type:ignore return [ "vars", - f'{",".join(t.bound_inputs)}', + f'{",".join(sorted(t.bound_inputs))}', "resolver", t.python_function_task.task_resolver.location, *t.python_function_task.task_resolver.loader_args(settings, t.python_function_task), diff --git a/flytekit/core/artifact.py b/flytekit/core/artifact.py index e32a924997..6c709f59a1 100644 --- a/flytekit/core/artifact.py +++ b/flytekit/core/artifact.py @@ -1,25 +1,13 @@ from __future__ import annotations import datetime -import random import typing from datetime import timedelta from typing import Optional, Union -from uuid import UUID -import isodate -from flyteidl.artifact import artifacts_pb2 from flyteidl.core import artifact_id_pb2 as art_id -from flyteidl.core.identifier_pb2 import TaskExecutionIdentifier, WorkflowExecutionIdentifier from google.protobuf.timestamp_pb2 import Timestamp -from flytekit.loggers import logger -from flytekit.models.literals import Literal -from flytekit.models.types import LiteralType - -if typing.TYPE_CHECKING: - from flytekit.remote.remote import FlyteRemote - TIME_PARTITION_KWARG = "time_partition" @@ -80,7 +68,7 @@ def bind_partitions(self, *args, **kwargs) -> ArtifactIDSpecification: p = Partitions(None) # k is the partition key, v should be static, or an input to the task or workflow for k, v in kwargs.items(): - if k not in self.artifact.partition_keys: + if not self.artifact.partition_keys or k not in self.artifact.partition_keys: raise ValueError(f"Partition key {k} not found in {self.artifact.partition_keys}") if isinstance(v, art_id.InputBindingData): p.partitions[k] = Partition(art_id.LabelValue(input_binding=v), name=k) @@ -95,9 +83,9 @@ def bind_partitions(self, *args, **kwargs) -> ArtifactIDSpecification: def to_partial_artifact_id(self) -> art_id.ArtifactID: # This function should only be called by transform_variable_map - artifact_id = self.artifact.to_flyte_idl().artifact_id + artifact_id = self.artifact.to_id_idl() # Use the partitions from this object, but replacement is not allowed by protobuf, so generate new object - p = partitions_to_idl(self.partitions) + p = Serializer.partitions_to_idl(self.partitions) tp = None if self.artifact.time_partitioned: if not self.time_partition: @@ -124,29 +112,6 @@ def to_partial_artifact_id(self) -> art_id.ArtifactID: return artifact_id -class ArtifactBindingData(object): - """ - We need some way of linking the triggered artifacts, with the parameter map. This object represents that link. - - These are used in two places in triggers. If the input to a trigger's target launchplan is the whole artifact, - then this binding should just have the index in the list. - If the input is an ArtifactQuery, then the query can reference one of these objects to dereference information - to be used in the search. - """ - - def __init__(self, triggered_artifact_id: int, partition_key: str, transform: Optional[str] = None): - self.triggered_artifact_id = triggered_artifact_id - self.partition_key = partition_key - self.transform = transform - - def to_flyte_idl(self) -> art_id.ArtifactBindingData: - return art_id.ArtifactBindingData( - index=self.triggered_artifact_id, - partition_key=self.partition_key, - transform=self.transform, - ) - - class ArtifactQuery(object): def __init__( self, @@ -161,12 +126,11 @@ def __init__( if not name: raise ValueError("Cannot create query without name") - # So normally, if you just do MyData.query(partitions={"region": "{{ inputs.region }}"}), it will just + # So normally, if you just do MyData.query(partitions="region": Inputs.region), it will just # use the input value to fill in the partition. But if you do - # MyData.query(partitions={"region": OtherArtifact.partitions.region }) + # MyData.query(region=OtherArtifact.partitions.region) # then you now have a dependency on the other artifact. This list keeps track of all the other Artifacts you've # referenced. - # Note that this is only used for Triggers. self.artifact = artifact bindings: typing.List[Artifact] = [] if time_partition: @@ -187,47 +151,9 @@ def __init__( def to_flyte_idl( self, - bindings: Optional[typing.List[Artifact]] = None, + **kwargs, ) -> art_id.ArtifactQuery: - """ - Think input keys can be removed - """ - - ak = art_id.ArtifactKey( - name=self.name, - project=self.project, - domain=self.domain, - ) - # If there's a tag, it takes precedence over other current options - if self.tag: - aq = art_id.ArtifactQuery( - artifact_tag=art_id.ArtifactTag(artifact_key=ak, value=art_id.LabelValue(static_value=self.tag)), - ) - return aq - - p = partitions_to_idl(self.partitions, bindings) - tp = None - if self.time_partition: - tp = self.time_partition.to_flyte_idl(bindings) - - i = art_id.ArtifactID( - artifact_key=ak, - partitions=p, - time_partition=tp, - ) - - aq = art_id.ArtifactQuery( - artifact_id=i, - ) - - return aq - - @staticmethod - def from_uri(uri: str) -> ArtifactQuery: - ... - - def as_uri(self) -> str: - ... + return Serializer.artifact_query_to_idl(self, **kwargs) class TimePartition(object): @@ -261,59 +187,19 @@ def __sub__(self, other: timedelta) -> TimePartition: tp.reference_artifact = self.reference_artifact return tp - def get_idl_partitions_for_trigger(self, bindings: typing.List[Artifact]) -> art_id.TimePartition: - if not self.reference_artifact or (self.reference_artifact and self.reference_artifact not in bindings): - # basically if there's no reference artifact, or if the reference artifact isn't - # in the list of triggers, then treat it like normal. - return art_id.TimePartition(value=self.value) - elif self.reference_artifact in bindings: - idx = bindings.index(self.reference_artifact) - transform = None - if self.op and self.other and isinstance(self.other, timedelta): - transform = str(self.op) + isodate.duration_isoformat(self.other) - lv = art_id.LabelValue( - triggered_binding=art_id.ArtifactBindingData( - index=idx, - bind_to_time_partition=True, - transform=transform, - ) - ) - return art_id.TimePartition(value=lv) - # investigate if this happens, if not, remove. else - logger.warning(f"Investigate - time partition in trigger with unhandled reference artifact {self}") - raise ValueError("Time partition reference artifact not found in ") - # return art_id.Partitions(value={TIME_PARTITION: self.value}) - - def to_flyte_idl(self, bindings: Optional[typing.List[Artifact]] = None) -> Optional[art_id.TimePartition]: - if bindings and len(bindings) > 0: - return self.get_idl_partitions_for_trigger(bindings) - - if not self.value: - # This is only for triggers - the backend needs to know of the existence of a time partition - return art_id.TimePartition() - - return art_id.TimePartition(value=self.value) - - -def partitions_to_idl( - partitions: Optional[Partitions], - bindings: Optional[typing.List[Artifact]] = None, -) -> Optional[art_id.Partitions]: - if partitions: - return partitions.to_flyte_idl(bindings) - - return None + def to_flyte_idl(self, **kwargs) -> Optional[art_id.TimePartition]: + return Serializer.time_partition_to_idl(self, **kwargs) class Partition(object): - def __init__(self, value: Optional[art_id.LabelValue], name: str = None): + def __init__(self, value: Optional[art_id.LabelValue], name: str): self.name = name self.value = value self.reference_artifact: Optional[Artifact] = None class Partitions(object): - def __init__(self, partitions: Optional[typing.Dict[str, Union[str, art_id.InputBindingData, Partition]]]): + def __init__(self, partitions: Optional[typing.Mapping[str, Union[str, art_id.InputBindingData, Partition]]]): self._partitions = {} if partitions: for k, v in partitions.items(): @@ -323,7 +209,7 @@ def __init__(self, partitions: Optional[typing.Dict[str, Union[str, art_id.Input self._partitions[k] = Partition(art_id.LabelValue(input_binding=v), name=k) else: self._partitions[k] = Partition(art_id.LabelValue(static_value=v), name=k) - self.reference_artifact = None + self.reference_artifact: Optional[Artifact] = None @property def partitions(self) -> Optional[typing.Dict[str, Partition]]: @@ -340,91 +226,22 @@ def __getattr__(self, item): return self.partitions[item] raise AttributeError(f"Partition {item} not found in {self}") - def get_idl_partitions_for_trigger( - self, - bindings: typing.List[Artifact] = None, - ) -> art_id.Partitions: - p = {} - # First create partition requirements for all the partitions - if self.reference_artifact and self.reference_artifact in bindings: - idx = bindings.index(self.reference_artifact) - triggering_artifact = bindings[idx] - if triggering_artifact.partition_keys: - for k in triggering_artifact.partition_keys: - p[k] = art_id.LabelValue( - triggered_binding=art_id.ArtifactBindingData( - index=idx, - partition_key=k, - ) - ) - - for k, v in self.partitions.items(): - if not v.reference_artifact or ( - v.reference_artifact - and v.reference_artifact is self.reference_artifact - and v.reference_artifact not in bindings - ): - # consider changing condition to just check for static value - p[k] = art_id.LabelValue(static_value=v.value.static_value) - elif v.reference_artifact in bindings: - # This line here is why the PartitionValue object has a name field. - # We might bind to a partition key that's a different name than the k here. - p[k] = art_id.LabelValue( - triggered_binding=art_id.ArtifactBindingData( - index=bindings.index(v.reference_artifact), - partition_key=v.name, - ) - ) - else: - raise ValueError(f"Partition has unhandled reference artifact {v.reference_artifact}") - - return art_id.Partitions(value=p) - - def to_flyte_idl( - self, - bindings: Optional[typing.List[Artifact]] = None, - ) -> Optional[art_id.Partitions]: - # This is basically a flag, which indicates that we are serializing this object within the context of a Trigger - # If we are not, then we are just serializing normally - if bindings and len(bindings) > 0: - return self.get_idl_partitions_for_trigger(bindings) - - if not self.partitions: - return None - - pp = {} - if self.partitions: - for k, v in self.partitions.items(): - if v.value is None: - # This should only happen when serializing for triggers - # Probably indicative of something in the data model that can be fixed - # down the road. - pp[k] = art_id.LabelValue(static_value="") - else: - pp[k] = v.value - return art_id.Partitions(value=pp) + def to_flyte_idl(self, **kwargs) -> Optional[art_id.Partitions]: + return Serializer.partitions_to_idl(self, **kwargs) class Artifact(object): """ An Artifact is effectively just a metadata layer on top of data that exists in Flyte. Most data of interest - will be the output of tasks and workflows. The other class is user uploads. - - This Python class has two purposes - as a Python representation of a materialized Artifact, - and as a way for users to specify that tasks/workflows create Artifacts and the manner - in which they are created. + will be the output of tasks and workflows. The other category is user uploads. - Python fields will be missing when retrieved from the service. - - Use one as input to workflow (only workflow for now) - df_artifact = Artifact.get("flyte://a1") - remote.execute(wf, inputs={"a": df_artifact}) + This Python class has limited purpose, as a way for users to specify that tasks/workflows create Artifacts + and the manner (i.e. name, partitions) in which they are created. Control creation parameters at task/workflow execution time :: @task - def t1() -> Annotated[nn.Module, Artifact(name="my.artifact.name", - tags=["latest", "1.0.0"])]: + def t1() -> Annotated[nn.Module, Artifact(name="my.artifact.name")]: ... """ @@ -438,16 +255,8 @@ def __init__( time_partition: Optional[TimePartition] = None, partition_keys: Optional[typing.List[str]] = None, partitions: Optional[Union[Partitions, typing.Dict[str, str]]] = None, - tags: Optional[typing.List[str]] = None, - python_val: Optional[typing.Any] = None, - python_type: Optional[typing.Type] = None, - literal: Optional[Literal] = None, - literal_type: Optional[LiteralType] = None, - short_description: Optional[str] = None, - source: Optional[typing.Union[WorkflowExecutionIdentifier, TaskExecutionIdentifier, str]] = None, ): """ - :param project: Should not be directly user provided, the project/domain will come from the project/domain of the execution that produced the output. These values will be filled in automatically when retrieving however. :param domain: See above. @@ -458,14 +267,6 @@ def __init__( :param partition_keys: This is a list of keys that will be used to partition the Artifact. These are not the values. Values are set via a () on the artifact and will end up in the partition_values field. :param partitions: This is a dictionary of partition keys to values. - :param tags: A list of tags that can be used as shortcuts to this Artifact. A tag targets one particular - partition if your Artifact is partitioned. - :param python_val: The Python value. - :param python_type: The Python type. - :param literal: The Literal value from the output. - :param literal_type: The LiteralType as taken from the task/workflow that produced the Artifact - :param short_description: A description of the Artifact. - TODO: Additional fields to come: figure out sources, and also add metadata (cards). """ if not name: raise ValueError("Can't instantiate an Artifact without a name.") @@ -479,27 +280,24 @@ def __init__( self._time_partition = time_partition self._time_partition.reference_artifact = self self.partition_keys = partition_keys - self._partitions = None + self._partitions: Optional[Partitions] = None if partitions: if isinstance(partitions, dict): self._partitions = Partitions(partitions) self.partition_keys = list(partitions.keys()) - else: + elif isinstance(partitions, Partitions): self._partitions = partitions + if not partitions.partitions: + raise ValueError("Partitions must be non-empty") self.partition_keys = list(partitions.partitions.keys()) + else: + raise ValueError(f"Partitions must be a dict or Partitions object, not {type(partitions)}") self._partitions.set_reference_artifact(self) if not partitions and partition_keys: # this should be the only time where we create Partition objects with None p = {k: Partition(None, name=k) for k in partition_keys} self._partitions = Partitions(p) self._partitions.set_reference_artifact(self) - self.python_val = python_val - self.python_type = python_type - self.literal = literal - self.literal_type = literal_type - self.tags = tags - self.short_description = short_description - self.source = source def __call__(self, *args, **kwargs) -> ArtifactIDSpecification: """ @@ -531,38 +329,17 @@ def __str__(self): f" name={self.name}\n" f" partitions={self.partitions}\n" f"{tp_str}" - f" tags={self.tags}\n" - f" literal_type=" - f"{self.literal_type}, " - f"literal={self.literal})" ) def __repr__(self): return self.__str__() - @classmethod - def get( - cls, - uri: Optional[str], - artifact_id: Optional[art_id.ArtifactID], - remote: FlyteRemote, - get_details: bool = False, - ) -> Optional[Artifact]: - """ - Use one locally. This retrieves the Literal. - a = remote.get("flyte://blah") - a = Artifact.get("flyte://blah", remote, tag="latest") - u = union.get("union://blah") - """ - return remote.get_artifact(uri=uri, artifact_id=artifact_id, get_details=get_details) - def query( self, project: Optional[str] = None, domain: Optional[str] = None, time_partition: Optional[Union[datetime.datetime, TimePartition, art_id.InputBindingData]] = None, partitions: Optional[Union[typing.Dict[str, str], Partitions]] = None, - tag: Optional[str] = None, **kwargs, ) -> ArtifactQuery: if self.partition_keys: @@ -576,13 +353,14 @@ def query( if partitions and kwargs: raise ValueError("Please either specify kwargs or a partitions object not both") + p_obj: Optional[Partitions] = None if kwargs: - partitions = Partitions(kwargs) - partitions.reference_artifact = self # only set top level + p_obj = Partitions(kwargs) + p_obj.reference_artifact = self # only set top level if partitions and isinstance(partitions, dict): - partitions = Partitions(partitions) - partitions.reference_artifact = self # only set top level + p_obj = Partitions(partitions) + p_obj.reference_artifact = self # only set top level tp = None if time_partition: @@ -600,59 +378,16 @@ def query( project=project or self.project or None, domain=domain or self.domain or None, time_partition=tp, - partitions=partitions or self.partitions, - tag=tag or self.tags[0] if self.tags else None, + partitions=p_obj or self.partitions, ) return aq - @classmethod - def initialize( - cls, - python_val: typing.Any, - python_type: typing.Type, - name: Optional[str] = None, - literal_type: Optional[LiteralType] = None, - tags: Optional[typing.List[str]] = None, - ) -> Artifact: - """ - Use this for when you have a Python value you want to get an Artifact object out of. - - This function readies an Artifact for creation, it doesn't actually create it just yet since this is a - network-less call. You will need to persist it with a FlyteRemote instance: - remote.create_artifact(Artifact.initialize(...)) - - Artifact.initialize("/path/to/file", tags={"tag1": "val1"}) - Artifact.initialize("/path/to/parquet", type=pd.DataFrame, tags=["0.1.0"]) - - What's set here is everything that isn't set by the server. What is set by the server? - - name, version, if not set by user. - - uri - Set by remote - - project, domain - """ - # Create the artifact object - return Artifact( - python_val=python_val, - python_type=python_type, - literal_type=literal_type, - tags=tags, - name=name, - ) - - # todo: merge this later with the as_artifact_id property - @property - def artifact_id(self) -> Optional[art_id.ArtifactID]: - if not self.project or not self.domain or not self.name or not self.version: - return None - - return self.to_flyte_idl().artifact_id - @property - def as_artifact_id(self) -> art_id.ArtifactID: + def concrete_artifact_id(self) -> art_id.ArtifactID: # This property is used when you want to ensure that this is a materialized artifact, all fields are known. if self.name is None or self.project is None or self.domain is None or self.version is None: raise ValueError("Cannot create artifact id without name, project, domain, version") - return self.to_flyte_idl().artifact_id + return self.to_id_idl() def embed_as_query( self, @@ -681,89 +416,100 @@ def embed_as_query( return aq - def to_flyte_idl(self) -> artifacts_pb2.Artifact: + def to_id_idl(self) -> art_id.ArtifactID: """ Converts this object to the IDL representation. This is here instead of translator because it's in the interface, a relatively simple proto object that's exposed to the user. """ - p = partitions_to_idl(self.partitions) - tp = self.time_partition.to_flyte_idl() if self.time_partitioned else None - - return artifacts_pb2.Artifact( - artifact_id=art_id.ArtifactID( - artifact_key=art_id.ArtifactKey( - project=self.project, - domain=self.domain, - name=self.name, - ), - version=self.version, - partitions=p, - time_partition=tp, + p = Serializer.partitions_to_idl(self.partitions) + tp = Serializer.time_partition_to_idl(self.time_partition) if self.time_partitioned else None + + i = art_id.ArtifactID( + artifact_key=art_id.ArtifactKey( + project=self.project, + domain=self.domain, + name=self.name, ), - spec=artifacts_pb2.ArtifactSpec(), - tags=self.tags, + version=self.version, + partitions=p, + time_partition=tp, ) - def as_create_request(self) -> artifacts_pb2.CreateArtifactRequest: - if not self.project or not self.domain: - raise ValueError("Project and domain are required to create an artifact") - name = self.name or UUID(int=random.getrandbits(128)).hex - ak = art_id.ArtifactKey(project=self.project, domain=self.domain, name=name) + return i + + +class ArtifactSerializationHandler(typing.Protocol): + """ + This protocol defines the interface for serializing artifact-related entities down to Flyte IDL. + """ + + def partitions_to_idl(self, p: Optional[Partitions], **kwargs) -> Optional[art_id.Partitions]: + ... + + def time_partition_to_idl(self, tp: Optional[TimePartition], **kwargs) -> Optional[art_id.TimePartition]: + ... + + def artifact_query_to_idl(self, aq: ArtifactQuery, **kwargs) -> art_id.ArtifactQuery: + ... + - spec = artifacts_pb2.ArtifactSpec( - value=self.literal, - type=self.literal_type, +class DefaultArtifactSerializationHandler(ArtifactSerializationHandler): + def partitions_to_idl(self, p: Optional[Partitions], **kwargs) -> Optional[art_id.Partitions]: + if p and p.partitions: + pp = {} + for k, v in p.partitions.items(): + if v.value is None: + # For specifying partitions in the Variable partial id + pp[k] = art_id.LabelValue(static_value="") + else: + pp[k] = v.value + return art_id.Partitions(value=pp) + return None + + def time_partition_to_idl(self, tp: Optional[TimePartition], **kwargs) -> Optional[art_id.TimePartition]: + if tp: + return art_id.TimePartition(value=tp.value) + return None + + def artifact_query_to_idl(self, aq: ArtifactQuery, **kwargs) -> art_id.ArtifactQuery: + ak = art_id.ArtifactKey( + name=aq.name, + project=aq.project, + domain=aq.domain, ) - partitions = partitions_to_idl(self.partitions) - tp = None - if self._time_partition: - tv = self.time_partition.value.time_value - if not tv: - raise Exception("missing time value") - tp = self.time_partition.value.time_value - - return artifacts_pb2.CreateArtifactRequest( - artifact_key=ak, spec=spec, partitions=partitions, time_partition_value=tp + p = self.partitions_to_idl(aq.partitions) + tp = self.time_partition_to_idl(aq.time_partition) + + i = art_id.ArtifactID( + artifact_key=ak, + partitions=p, + time_partition=tp, ) - @classmethod - def from_flyte_idl(cls, pb2: artifacts_pb2.Artifact) -> Artifact: - """ - Converts the IDL representation to this object. - """ - tags = [t for t in pb2.tags] if pb2.tags else None - a = Artifact( - project=pb2.artifact_id.artifact_key.project, - domain=pb2.artifact_id.artifact_key.domain, - name=pb2.artifact_id.artifact_key.name, - version=pb2.artifact_id.version, - tags=tags, - literal_type=LiteralType.from_flyte_idl(pb2.spec.type), - literal=Literal.from_flyte_idl(pb2.spec.value), - # source=pb2.spec.source, # todo: source isn't installed in artifact service yet + aq = art_id.ArtifactQuery( + artifact_id=i, ) - if pb2.artifact_id.HasField("partitions"): - if len(pb2.artifact_id.partitions.value) > 0: - # static values should be the only ones set since currently we don't from_flyte_idl - # anything that's not a materialized artifact. - # if TIME_PARTITION in pb2.artifact_id.partitions.value: - # a._time_partition = TimePartition(pb2.artifact_id.partitions.value[TIME_PARTITION].static_value) - # a._time_partition.reference_artifact = a - - a._partitions = Partitions( - partitions={ - k: Partition(value=v, name=k) - for k, v in pb2.artifact_id.partitions.value.items() - # if k != TIME_PARTITION - } - ) - a.partitions.reference_artifact = a - if pb2.artifact_id.HasField("time_partition"): - ts = pb2.artifact_id.time_partition.value.time_value - dt = ts.ToDatetime() - a._time_partition = TimePartition(dt) - a._time_partition.reference_artifact = a - - return a + + return aq + + +class Serializer(object): + serializer: ArtifactSerializationHandler = DefaultArtifactSerializationHandler() + + @classmethod + def register_serializer(cls, serializer: ArtifactSerializationHandler): + cls.serializer = serializer + + @classmethod + def partitions_to_idl(cls, p: Optional[Partitions], **kwargs) -> Optional[art_id.Partitions]: + return cls.serializer.partitions_to_idl(p, **kwargs) + + @classmethod + def time_partition_to_idl(cls, tp: TimePartition, **kwargs) -> Optional[art_id.TimePartition]: + return cls.serializer.time_partition_to_idl(tp, **kwargs) + + @classmethod + def artifact_query_to_idl(cls, aq: ArtifactQuery, **kwargs) -> art_id.ArtifactQuery: + return cls.serializer.artifact_query_to_idl(aq, **kwargs) diff --git a/flytekit/core/base_task.py b/flytekit/core/base_task.py index b9171f54b8..61b2ca89d4 100644 --- a/flytekit/core/base_task.py +++ b/flytekit/core/base_task.py @@ -27,7 +27,7 @@ from flyteidl.core import tasks_pb2 -from flytekit.configuration import SerializationSettings +from flytekit.configuration import LocalConfig, SerializationSettings from flytekit.core.context_manager import ( ExecutionParameters, ExecutionState, @@ -265,7 +265,8 @@ def local_execute( input_literal_map = _literal_models.LiteralMap(literals=kwargs) # if metadata.cache is set, check memoized version - if self.metadata.cache: + local_config = LocalConfig.auto() + if self.metadata.cache and local_config.cache_enabled: # TODO: how to get a nice `native_inputs` here? logger.info( f"Checking cache for task named {self.name}, cache version {self.metadata.cache_version} " diff --git a/flytekit/core/data_persistence.py b/flytekit/core/data_persistence.py index 1579ef3f6b..b597c75b56 100644 --- a/flytekit/core/data_persistence.py +++ b/flytekit/core/data_persistence.py @@ -247,7 +247,7 @@ def get(self, from_path: str, to_path: str, recursive: bool = False, **kwargs): return shutil.copytree( self.strip_file_header(from_path), self.strip_file_header(to_path), dirs_exist_ok=True ) - print(f"Getting {from_path} to {to_path}") + logger.info(f"Getting {from_path} to {to_path}") dst = file_system.get(from_path, to_path, recursive=recursive, **kwargs) if isinstance(dst, (str, pathlib.Path)): return dst diff --git a/flytekit/core/interface.py b/flytekit/core/interface.py index b5b2e9eb0e..99add602b1 100644 --- a/flytekit/core/interface.py +++ b/flytekit/core/interface.py @@ -224,7 +224,7 @@ def transform_inputs_to_parameters( if isinstance(_default, ArtifactQuery): params[k] = _interface_models.Parameter(var=v, required=False, artifact_query=_default.to_flyte_idl()) elif isinstance(_default, Artifact): - artifact_id = _default.as_artifact_id # may raise + artifact_id = _default.concrete_artifact_id # may raise params[k] = _interface_models.Parameter(var=v, required=False, artifact_id=artifact_id) else: required = _default is None @@ -355,7 +355,7 @@ def transform_function_to_interface(fn: typing.Callable, docstring: Optional[Doc def transform_variable_map( variable_map: Dict[str, type], - descriptions: Dict[str, str] = None, + descriptions: Optional[Dict[str, str]] = None, ) -> Dict[str, _interface_models.Variable]: """ Given a map of str (names of inputs for instance) to their Python native types, return a map of the name to a @@ -370,7 +370,7 @@ def transform_variable_map( def detect_artifact( - ts: typing.Tuple[typing.Any], + ts: typing.Tuple[typing.Any, ...], ) -> Optional[art_id.ArtifactID]: """ If the user wishes to control how Artifacts are created (i.e. naming them, etc.) this is where we pick it up and diff --git a/flytekit/core/map_task.py b/flytekit/core/map_task.py index 1201a3ede0..aac31a1ee9 100644 --- a/flytekit/core/map_task.py +++ b/flytekit/core/map_task.py @@ -92,7 +92,8 @@ def __init__( f = actual_task.lhs else: _, mod, f, _ = tracker.extract_task_module(typing.cast(PythonFunctionTask, actual_task).task_function) - h = hashlib.md5(collection_interface.__str__().encode("utf-8")).hexdigest() + sorted_bounded_inputs = ",".join(sorted(self._bound_inputs)) + h = hashlib.md5(sorted_bounded_inputs.encode("utf-8")).hexdigest() name = f"{mod}.map_{f}_{h}" self._cmd_prefix: typing.Optional[typing.List[str]] = None @@ -404,7 +405,7 @@ def load_task(self, loader_args: List[str], max_concurrency: int = 0) -> MapPyth def loader_args(self, settings: SerializationSettings, t: MapPythonTask) -> List[str]: # type:ignore return [ "vars", - f'{",".join(t.bound_inputs)}', + f'{",".join(sorted(t.bound_inputs))}', "resolver", t.run_task.task_resolver.location, *t.run_task.task_resolver.loader_args(settings, t.run_task), diff --git a/flytekit/core/node.py b/flytekit/core/node.py index 6c84877c90..b7ab8692b6 100644 --- a/flytekit/core/node.py +++ b/flytekit/core/node.py @@ -145,6 +145,15 @@ def with_overrides(self, *args, **kwargs): limits = kwargs.get("limits") if limits and not isinstance(limits, Resources): raise AssertionError("limits should be specified as flytekit.Resources") + + if not limits: + logger.warning( + ( + f"Requests overridden on node {self.id} ({self.metadata.short_string()}) without specifying limits. " + "Requests are clamped to original limits." + ) + ) + resources = convert_resources_to_resource_model(requests=requests, limits=limits) assert_no_promises_in_resources(resources) self._resources = resources diff --git a/flytekit/core/python_function_task.py b/flytekit/core/python_function_task.py index e1e80a4227..c26bdd6c6e 100644 --- a/flytekit/core/python_function_task.py +++ b/flytekit/core/python_function_task.py @@ -13,11 +13,14 @@ """ +from __future__ import annotations + from abc import ABC from collections import OrderedDict from enum import Enum -from typing import Any, Callable, List, Optional, TypeVar, Union, cast +from typing import Any, Callable, Iterable, List, Optional, TypeVar, Union, cast +from flytekit.core import launch_plan as _annotated_launch_plan from flytekit.core.base_task import Task, TaskResolverMixin from flytekit.core.context_manager import ExecutionState, FlyteContext, FlyteContextManager from flytekit.core.docstring import Docstring @@ -27,6 +30,7 @@ from flytekit.core.tracker import extract_task_module, is_functools_wrapped_module_level, isnested, istestfunction from flytekit.core.workflow import ( PythonFunctionWorkflow, + WorkflowBase, WorkflowFailurePolicy, WorkflowMetadata, WorkflowMetadataDefaults, @@ -102,6 +106,9 @@ def __init__( ignore_input_vars: Optional[List[str]] = None, execution_mode: ExecutionBehavior = ExecutionBehavior.DEFAULT, task_resolver: Optional[TaskResolverMixin] = None, + node_dependency_hints: Optional[ + Iterable[Union["PythonFunctionTask", "_annotated_launch_plan.LaunchPlan", WorkflowBase]] + ] = None, **kwargs, ): """ @@ -112,6 +119,9 @@ def __init__( :param Optional[ExecutionBehavior] execution_mode: Defines how the execution should behave, for example executing normally or specially handling a dynamic case. :param str task_type: String task type to be associated with this Task + :param Optional[Iterable[Union["PythonFunctionTask", "_annotated_launch_plan.LaunchPlan", WorkflowBase]]] node_dependency_hints: + A list of tasks, launchplans, or workflows that this task depends on. This is only + for dynamic tasks/workflows, where flyte cannot automatically determine the dependencies prior to runtime. """ if task_function is None: raise ValueError("TaskFunction is a required parameter for PythonFunctionTask") @@ -145,12 +155,24 @@ def __init__( ) self._task_function = task_function self._execution_mode = execution_mode + self._node_dependency_hints = node_dependency_hints + if self._node_dependency_hints is not None and self._execution_mode != self.ExecutionBehavior.DYNAMIC: + raise ValueError( + "node_dependency_hints should only be used on dynamic tasks. On static tasks and " + "workflows its redundant because flyte can find the node dependencies automatically" + ) self._wf = None # For dynamic tasks @property def execution_mode(self) -> ExecutionBehavior: return self._execution_mode + @property + def node_dependency_hints( + self, + ) -> Optional[Iterable[Union["PythonFunctionTask", "_annotated_launch_plan.LaunchPlan", WorkflowBase]]]: + return self._node_dependency_hints + @property def task_function(self): return self._task_function diff --git a/flytekit/core/task.py b/flytekit/core/task.py index 547abd41fa..a99fbf599e 100644 --- a/flytekit/core/task.py +++ b/flytekit/core/task.py @@ -1,7 +1,11 @@ +from __future__ import annotations + import datetime as _datetime from functools import update_wrapper -from typing import Any, Callable, Dict, List, Optional, Type, TypeVar, Union, overload +from typing import Any, Callable, Dict, Iterable, List, Optional, Type, TypeVar, Union, overload +from flytekit.core import launch_plan as _annotated_launchplan +from flytekit.core import workflow as _annotated_workflow from flytekit.core.base_task import TaskMetadata, TaskResolverMixin from flytekit.core.interface import transform_function_to_interface from flytekit.core.pod_template import PodTemplate @@ -97,6 +101,9 @@ def task( limits: Optional[Resources] = ..., secret_requests: Optional[List[Secret]] = ..., execution_mode: PythonFunctionTask.ExecutionBehavior = ..., + node_dependency_hints: Optional[ + Iterable[Union[PythonFunctionTask, _annotated_launchplan.LaunchPlan, _annotated_workflow.WorkflowBase]] + ] = ..., task_resolver: Optional[TaskResolverMixin] = ..., docs: Optional[Documentation] = ..., disable_deck: Optional[bool] = ..., @@ -125,6 +132,9 @@ def task( limits: Optional[Resources] = ..., secret_requests: Optional[List[Secret]] = ..., execution_mode: PythonFunctionTask.ExecutionBehavior = ..., + node_dependency_hints: Optional[ + Iterable[Union[PythonFunctionTask, _annotated_launchplan.LaunchPlan, _annotated_workflow.WorkflowBase]] + ] = ..., task_resolver: Optional[TaskResolverMixin] = ..., docs: Optional[Documentation] = ..., disable_deck: Optional[bool] = ..., @@ -152,6 +162,9 @@ def task( limits: Optional[Resources] = None, secret_requests: Optional[List[Secret]] = None, execution_mode: PythonFunctionTask.ExecutionBehavior = PythonFunctionTask.ExecutionBehavior.DEFAULT, + node_dependency_hints: Optional[ + Iterable[Union[PythonFunctionTask, _annotated_launchplan.LaunchPlan, _annotated_workflow.WorkflowBase]] + ] = None, task_resolver: Optional[TaskResolverMixin] = None, docs: Optional[Documentation] = None, disable_deck: Optional[bool] = None, @@ -246,6 +259,28 @@ def foo2(): Refer to :py:class:`Secret` to understand how to specify the request for a secret. It may change based on the backend provider. :param execution_mode: This is mainly for internal use. Please ignore. It is filled in automatically. + :param node_dependency_hints: A list of tasks, launchplans, or workflows that this task depends on. This is only + for dynamic tasks/workflows, where flyte cannot automatically determine the dependencies prior to runtime. + Even on dynamic tasks this is optional, but in some scenarios it will make registering the workflow easier, + because it allows registration to be done the same as for static tasks/workflows. + + For example this is useful to run launchplans dynamically, because launchplans must be registered on flyteadmin + before they can be run. Tasks and workflows do not have this requirement. + + .. code-block:: python + + @workflow + def workflow0(): + ... + + launchplan0 = LaunchPlan.get_or_create(workflow0) + + # Specify node_dependency_hints so that launchplan0 will be registered on flyteadmin, despite this being a + # dynamic task. + @dynamic(node_dependency_hints=[launchplan0]) + def launch_dynamically(): + # To run a sub-launchplan it must have previously been registered on flyteadmin. + return [launchplan0]*10 :param task_resolver: Provide a custom task resolver. :param disable_deck: (deprecated) If true, this task will not output deck html file :param enable_deck: If true, this task will output deck html file @@ -276,6 +311,7 @@ def wrapper(fn: Callable[..., Any]) -> PythonFunctionTask[T]: limits=limits, secret_requests=secret_requests, execution_mode=execution_mode, + node_dependency_hints=node_dependency_hints, task_resolver=task_resolver, disable_deck=disable_deck, enable_deck=enable_deck, diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 4220d6e7b6..76f750233b 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -325,6 +325,13 @@ class Test(DataClassJsonMixin): def __init__(self): super().__init__("Object-Dataclass-Transformer", object) + self._serializable_classes = [DataClassJSONMixin, DataClassJsonMixin] + try: + from mashumaro.mixins.orjson import DataClassORJSONMixin + + self._serializable_classes.append(DataClassORJSONMixin) + except ModuleNotFoundError: + pass def assert_type(self, expected_type: Type[DataClassJsonMixin], v: T): # Skip iterating all attributes in the dataclass if the type of v already matches the expected_type @@ -417,7 +424,7 @@ def get_literal_type(self, t: Type[T]) -> LiteralType: f"Type {t} cannot be parsed." ) - if not issubclass(t, DataClassJsonMixin) and not issubclass(t, DataClassJSONMixin): + if not self.is_serializable_class(t): raise AssertionError( f"Dataclass {t} should be decorated with @dataclass_json or mixin with DataClassJSONMixin to be " f"serialized correctly" @@ -465,6 +472,9 @@ def get_literal_type(self, t: Type[T]) -> LiteralType: return _type_models.LiteralType(simple=_type_models.SimpleType.STRUCT, metadata=schema, structure=ts) + def is_serializable_class(self, class_: Type[T]) -> bool: + return any(issubclass(class_, serializable_class) for serializable_class in self._serializable_classes) + def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], expected: LiteralType) -> Literal: if isinstance(python_val, dict): json_str = json.dumps(python_val) @@ -475,9 +485,7 @@ def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], exp f"{type(python_val)} is not of type @dataclass, only Dataclasses are supported for " f"user defined datatypes in Flytekit" ) - if not issubclass(type(python_val), DataClassJsonMixin) and not issubclass( - type(python_val), DataClassJSONMixin - ): + if not self.is_serializable_class(type(python_val)): raise TypeTransformerFailedError( f"Dataclass {python_type} should be decorated with @dataclass_json or inherit DataClassJSONMixin to be " f"serialized correctly" @@ -730,9 +738,7 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: f"{expected_python_type} is not of type @dataclass, only Dataclasses are supported for " "user defined datatypes in Flytekit" ) - if not issubclass(expected_python_type, DataClassJsonMixin) and not issubclass( - expected_python_type, DataClassJSONMixin - ): + if not self.is_serializable_class(expected_python_type): raise TypeTransformerFailedError( f"Dataclass {expected_python_type} should be decorated with @dataclass_json or mixin with DataClassJSONMixin to be " f"serialized correctly" diff --git a/flytekit/core/workflow.py b/flytekit/core/workflow.py index 7a54ef6de5..108b323a48 100644 --- a/flytekit/core/workflow.py +++ b/flytekit/core/workflow.py @@ -9,6 +9,7 @@ from typing import Any, Callable, Coroutine, Dict, List, Optional, Tuple, Type, Union, cast, overload from flytekit.core import constants as _common_constants +from flytekit.core import launch_plan as _annotated_launch_plan from flytekit.core.base_task import PythonTask, Task from flytekit.core.class_based_resolver import ClassStorageTaskResolver from flytekit.core.condition import ConditionalSection, conditional @@ -20,8 +21,11 @@ FlyteEntities, ) from flytekit.core.docstring import Docstring -from flytekit.core.interface import Interface, transform_function_to_interface, transform_interface_to_typed_interface -from flytekit.core.launch_plan import LaunchPlan +from flytekit.core.interface import ( + Interface, + transform_function_to_interface, + transform_interface_to_typed_interface, +) from flytekit.core.node import Node from flytekit.core.promise import ( NodeOutput, @@ -523,7 +527,7 @@ def create_conditional(self, name: str) -> ConditionalSection: FlyteContextManager.with_context(ctx.with_compilation_state(self.compilation_state)) return conditional(name=name) - def add_entity(self, entity: Union[PythonTask, LaunchPlan, WorkflowBase], **kwargs) -> Node: + def add_entity(self, entity: Union[PythonTask, _annotated_launch_plan.LaunchPlan, WorkflowBase], **kwargs) -> Node: """ Anytime you add an entity, all the inputs to the entity must be bound. """ @@ -606,7 +610,7 @@ def add_workflow_output( def add_task(self, task: PythonTask, **kwargs) -> Node: return self.add_entity(task, **kwargs) - def add_launch_plan(self, launch_plan: LaunchPlan, **kwargs) -> Node: + def add_launch_plan(self, launch_plan: _annotated_launch_plan.LaunchPlan, **kwargs) -> Node: return self.add_entity(launch_plan, **kwargs) def add_subwf(self, sub_wf: WorkflowBase, **kwargs) -> Node: @@ -739,14 +743,19 @@ def compile(self, **kwargs): ) workflow_outputs = workflow_outputs[0] t = self.python_interface.outputs[output_names[0]] - b, _ = binding_from_python_std( - ctx, - output_names[0], - self.interface.outputs[output_names[0]].type, - workflow_outputs, - t, - ) - bindings.append(b) + try: + b, _ = binding_from_python_std( + ctx, + output_names[0], + self.interface.outputs[output_names[0]].type, + workflow_outputs, + t, + ) + bindings.append(b) + except Exception as e: + raise FlyteValidationException( + f"Failed to bind output {output_names[0]} for function {self.name}: {e}" + ) from e elif len(output_names) > 1: if not isinstance(workflow_outputs, tuple): raise AssertionError("The Workflow specification indicates multiple return values, received only one") @@ -756,14 +765,17 @@ def compile(self, **kwargs): if isinstance(workflow_outputs[i], ConditionalSection): raise AssertionError("A Conditional block (if-else) should always end with an `else_()` clause") t = self.python_interface.outputs[out] - b, _ = binding_from_python_std( - ctx, - out, - self.interface.outputs[out].type, - workflow_outputs[i], - t, - ) - bindings.append(b) + try: + b, _ = binding_from_python_std( + ctx, + out, + self.interface.outputs[out].type, + workflow_outputs[i], + t, + ) + bindings.append(b) + except Exception as e: + raise FlyteValidationException(f"Failed to bind output {out} for function {self.name}: {e}") from e # Save all the things necessary to create an WorkflowTemplate, except for the missing project and domain self._nodes = all_nodes diff --git a/flytekit/exceptions/scopes.py b/flytekit/exceptions/scopes.py index f0a0f02189..a9a33b748d 100644 --- a/flytekit/exceptions/scopes.py +++ b/flytekit/exceptions/scopes.py @@ -39,7 +39,7 @@ def verbose_message(self): traceback_str = "\n ".join([""] + lines) format_str = "Traceback (most recent call last):\n" "{traceback}\n" "\n" "Message:\n" "\n" " {message}" - return format_str.format(traceback=traceback_str, message=str(self.value)) + return format_str.format(traceback=traceback_str, message=f"{self.type.__name__}: {self.value}") def __str__(self): return str(self.value) diff --git a/flytekit/extend/backend/agent_service.py b/flytekit/extend/backend/agent_service.py index 73737f3a6c..588631321d 100644 --- a/flytekit/extend/backend/agent_service.py +++ b/flytekit/extend/backend/agent_service.py @@ -7,10 +7,14 @@ CreateTaskResponse, DeleteTaskRequest, DeleteTaskResponse, + GetAgentRequest, + GetAgentResponse, GetTaskRequest, GetTaskResponse, + ListAgentsRequest, + ListAgentsResponse, ) -from flyteidl.service.agent_pb2_grpc import AsyncAgentServiceServicer +from flyteidl.service.agent_pb2_grpc import AgentMetadataServiceServicer, AsyncAgentServiceServicer from prometheus_client import Counter, Summary from flytekit import logger @@ -26,18 +30,20 @@ # Follow the naming convention. https://prometheus.io/docs/practices/naming/ request_success_count = Counter( - f"{metric_prefix}requests_success_total", "Total number of successful requests", ["task_type", "operation"] + f"{metric_prefix}requests_success_total", + "Total number of successful requests", + ["task_type", "operation"], ) request_failure_count = Counter( f"{metric_prefix}requests_failure_total", "Total number of failed requests", ["task_type", "operation", "error_code"], ) - request_latency = Summary( - f"{metric_prefix}request_latency_seconds", "Time spent processing agent request", ["task_type", "operation"] + f"{metric_prefix}request_latency_seconds", + "Time spent processing agent request", + ["task_type", "operation"], ) - input_literal_size = Summary(f"{metric_prefix}input_literal_bytes", "Size of input literal", ["task_type"]) @@ -96,8 +102,12 @@ async def CreateTask(self, request: CreateTaskRequest, context: grpc.ServicerCon logger.info(f"{tmp.type} agent start creating the job") if agent.asynchronous: return await agent.async_create( - context=context, inputs=inputs, output_prefix=request.output_prefix, task_template=tmp + context=context, + inputs=inputs, + output_prefix=request.output_prefix, + task_template=tmp, ) + return await asyncio.get_running_loop().run_in_executor( None, agent.create, @@ -122,3 +132,12 @@ async def DeleteTask(self, request: DeleteTaskRequest, context: grpc.ServicerCon if agent.asynchronous: return await agent.async_delete(context=context, resource_meta=request.resource_meta) return await asyncio.get_running_loop().run_in_executor(None, agent.delete, context, request.resource_meta) + + +class AgentMetadataService(AgentMetadataServiceServicer): + async def GetAgent(self, request: GetAgentRequest, context: grpc.ServicerContext) -> GetAgentResponse: + return GetAgentResponse(agent=AgentRegistry._METADATA[request.name]) + + async def ListAgents(self, request: ListAgentsRequest, context: grpc.ServicerContext) -> ListAgentsResponse: + agents = [agent for agent in AgentRegistry._METADATA.values()] + return ListAgentsResponse(agents=agents) diff --git a/flytekit/extend/backend/base_agent.py b/flytekit/extend/backend/base_agent.py index eb8e7562b6..28bc18ddd2 100644 --- a/flytekit/extend/backend/base_agent.py +++ b/flytekit/extend/backend/base_agent.py @@ -14,6 +14,7 @@ RETRYABLE_FAILURE, RUNNING, SUCCEEDED, + Agent, CreateTaskResponse, DeleteTaskResponse, GetTaskResponse, @@ -45,7 +46,9 @@ class AgentBase(ABC): will look up the agent based on the task type. Every task type can only have one agent. """ - def __init__(self, task_type: str, asynchronous=True): + name = "Base Agent" + + def __init__(self, task_type: str, asynchronous: bool = True): self._task_type = task_type self._asynchronous = asynchronous @@ -113,18 +116,28 @@ async def async_delete(self, context: grpc.ServicerContext, resource_meta: bytes class AgentRegistry(object): """ - This is the registry for all agents. The agent service will look up the agent - based on the task type. + This is the registry for all agents. + The agent service will look up the agent registry based on the task type. + The agent metadata service will look up the agent metadata based on the agent name. """ _REGISTRY: typing.Dict[str, AgentBase] = {} + _METADATA: typing.Dict[str, Agent] = {} @staticmethod def register(agent: AgentBase): if agent.task_type in AgentRegistry._REGISTRY: raise ValueError(f"Duplicate agent for task type {agent.task_type}") AgentRegistry._REGISTRY[agent.task_type] = agent - logger.info(f"Registering an agent for task type {agent.task_type}") + + if agent.name in AgentRegistry._METADATA: + agent_metadata = AgentRegistry._METADATA[agent.name] + agent_metadata.supported_task_types.append(agent.task_type) + else: + agent_metadata = Agent(name=agent.name, supported_task_types=[agent.task_type]) + AgentRegistry._METADATA[agent.name] = agent_metadata + + logger.info(f"Registering an agent for task type: {agent.task_type}, name: {agent.name}") @staticmethod def get_agent(task_type: str) -> typing.Optional[AgentBase]: @@ -132,6 +145,12 @@ def get_agent(task_type: str) -> typing.Optional[AgentBase]: raise FlyteAgentNotFound(f"Cannot find agent for task type: {task_type}.") return AgentRegistry._REGISTRY[task_type] + @staticmethod + def get_agent_metadata(name: str) -> Agent: + if name not in AgentRegistry._METADATA: + raise FlyteAgentNotFound(f"Cannot find agent for name: {name}.") + return AgentRegistry._METADATA[name] + def convert_to_flyte_state(state: str) -> State: """ diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index cef455acbb..7a8ef547da 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -7,10 +7,14 @@ from abc import abstractmethod from dataclasses import asdict, dataclass from functools import lru_cache -from typing import List, Optional, Union +from importlib import metadata +from typing import Dict, List, Optional, Tuple, Union import click import requests +from packaging.version import Version + +from flytekit.exceptions.user import FlyteAssertion DOCKER_HUB = "docker.io" _F_IMG_ID = "_F_IMG_ID" @@ -44,7 +48,7 @@ class ImageSpec: name: str = "flytekit" python_version: str = None # Use default python in the base image if None. - builder: str = "envd" + builder: Optional[str] = None source_root: Optional[str] = None env: Optional[typing.Dict[str, str]] = None registry: Optional[str] = None @@ -67,9 +71,15 @@ def __post_init__(self): self.registry = self.registry.lower() def image_name(self) -> str: - """ - return full image name with tag. - """ + """Full image name with tag.""" + image_name = self._image_name() + try: + return ImageBuildEngine._IMAGE_NAME_TO_REAL_NAME[image_name] + except KeyError: + return image_name + + def _image_name(self) -> str: + """Construct full image name with tag.""" tag = calculate_hash_from_image_spec(self) container_image = f"{self.name}:{tag}" if self.registry: @@ -174,12 +184,15 @@ def with_apt_packages(self, apt_packages: Union[str, List[str]]) -> "ImageSpec": class ImageSpecBuilder: @abstractmethod - def build_image(self, image_spec: ImageSpec): + def build_image(self, image_spec: ImageSpec) -> Optional[str]: """ Build the docker image and push it to the registry. Args: image_spec: image spec of the task. + + Returns: + fully_qualified_image_name: Fully qualified image name. If None, then `image_spec.image_name()` is used. """ raise NotImplementedError("This method is not implemented in the base class.") @@ -189,24 +202,45 @@ class ImageBuildEngine: ImageBuildEngine contains a list of builders that can be used to build an ImageSpec. """ - _REGISTRY: typing.Dict[str, ImageSpecBuilder] = {} + _REGISTRY: typing.Dict[str, Tuple[ImageSpecBuilder, int]] = {} _BUILT_IMAGES: typing.Set[str] = set() + # _IMAGE_NAME_TO_REAL_NAME is used to keep track of the fully qualified image name + # returned by the image builder. This allows ImageSpec to map from `image_spc.image_name()` + # to the real qualified name. + _IMAGE_NAME_TO_REAL_NAME: Dict[str, str] = {} @classmethod - def register(cls, builder_type: str, image_spec_builder: ImageSpecBuilder): - cls._REGISTRY[builder_type] = image_spec_builder + def register(cls, builder_type: str, image_spec_builder: ImageSpecBuilder, priority: int = 5): + cls._REGISTRY[builder_type] = (image_spec_builder, priority) @classmethod @lru_cache - def build(cls, image_spec: ImageSpec): + def build(cls, image_spec: ImageSpec) -> str: + if image_spec.builder is None and cls._REGISTRY: + builder = max(cls._REGISTRY, key=lambda name: cls._REGISTRY[name][1]) + else: + builder = image_spec.builder + img_name = image_spec.image_name() if img_name in cls._BUILT_IMAGES or image_spec.exist(): click.secho(f"Image {img_name} found. Skip building.", fg="blue") else: click.secho(f"Image {img_name} not found. Building...", fg="blue") - if image_spec.builder not in cls._REGISTRY: - raise Exception(f"Builder {image_spec.builder} is not registered.") - cls._REGISTRY[image_spec.builder].build_image(image_spec) + if builder not in cls._REGISTRY: + raise Exception(f"Builder {builder} is not registered.") + if builder == "envd": + envd_version = metadata.version("envd") + # flytekit v1.10.2+ copies the workflow code to the WorkDir specified in the Dockerfile. However, envd<0.3.39 + # overwrites the WorkDir when building the image, resulting in a permission issue when flytekit downloads the file. + if Version(envd_version) < Version("0.3.39"): + raise FlyteAssertion( + f"envd version {envd_version} is not compatible with flytekit>v1.10.2." + f" Please upgrade envd to v0.3.39+." + ) + + fully_qualified_image_name = cls._REGISTRY[builder][0].build_image(image_spec) + if fully_qualified_image_name is not None: + cls._IMAGE_NAME_TO_REAL_NAME[img_name] = fully_qualified_image_name cls._BUILT_IMAGES.add(img_name) diff --git a/flytekit/remote/remote.py b/flytekit/remote/remote.py index 46c70bd49d..c8acaea259 100644 --- a/flytekit/remote/remote.py +++ b/flytekit/remote/remote.py @@ -24,8 +24,6 @@ import fsspec import requests from flyteidl.admin.signal_pb2 import Signal, SignalListRequest, SignalSetRequest -from flyteidl.artifact import artifacts_pb2 -from flyteidl.core import artifact_id_pb2 as art_id from flyteidl.core import literals_pb2 from flytekit.clients.friendly import SynchronousFlyteClient @@ -362,72 +360,6 @@ def _fetch(): return LazyEntity(name=name, getter=_fetch) - def create_artifact(self, artifact: Artifact): - """ - Create an artifact in FlyteAdmin. - - :param artifact: The artifact to create. - :return: The artifact as persisted in the service. - """ - # Two things can happen here - - # - the call to to_literal may upload something, in the case of an offloaded data type. - # - if this happens, the upload request should already return the created Artifact object. - if artifact.literal is None: - with self.remote_context() as ctx: - lt = artifact.literal_type or TypeEngine.to_literal_type(artifact.python_type) - lit = TypeEngine.to_literal(ctx, artifact.python_val, artifact.python_type, lt) - artifact.literal_type = lt.to_flyte_idl() - artifact.literal = lit.to_flyte_idl() - else: - raise ValueError("Cannot create an artifact with a literal already set.") - - # Need to detect here what happened. If the conversion triggered a data upload, then the information from - # Artifact object from that call should be copied in. - # If not, an explicit call to create_artifact should be made. - - if artifact.project is None: - artifact.project = self.default_project - if artifact.domain is None: - artifact.domain = self.default_domain - resp = self.client.create_artifact(artifact.as_create_request()) - print(f"Res is {resp}") - artifact.project = resp.artifact.artifact_id.artifact_key.project - artifact.domain = resp.artifact.artifact_id.artifact_key.domain - artifact.name = resp.artifact.artifact_id.artifact_key.name - artifact.version = resp.artifact.artifact_id.version - artifact.source = ( - resp.artifact.spec.principal or resp.artifact.spec.execution or resp.artifact.spec.task_execution - ) - - def get_artifact( - self, - uri: typing.Optional[str] = None, - artifact_key: typing.Optional[art_id.ArtifactKey] = None, - artifact_id: typing.Optional[art_id.ArtifactID] = None, - query: typing.Optional[art_id.ArtifactQuery] = None, - tag: typing.Optional[str] = None, - get_details: bool = False, - ) -> typing.Optional[Artifact]: - if query: - q = query - elif uri: - q = art_id.ArtifactQuery(uri=uri) - elif artifact_key: - if tag: - q = art_id.ArtifactQuery(artifact_tag=art_id.ArtifactTag(artifact_key=artifact_key, tag=tag)) - else: - q = art_id.ArtifactQuery(artifact_id=art_id.ArtifactID(artifact_key=artifact_key)) - elif artifact_id: - if tag: - raise ValueError("If using tag specify key instead of ID.") - q = art_id.ArtifactQuery(artifact_id=artifact_id) - else: - raise ValueError("One of uri, key, id") - req = artifacts_pb2.GetArtifactRequest(query=q, details=get_details) - resp = self.client.get_artifact(req) - a = Artifact.from_flyte_idl(resp.artifact) - return a - def fetch_workflow( self, project: str = None, domain: str = None, name: str = None, version: str = None ) -> FlyteWorkflow: @@ -949,7 +881,9 @@ def _version_from_hash( for s in additional_context: h.update(bytes(s, "utf-8")) - return base64.urlsafe_b64encode(h.digest()).decode("ascii") + # Omit the character '=' from the version as that's essentially padding used by the base64 encoding + # and does not increase entropy of the hash while making it very inconvenient to copy-and-paste. + return base64.urlsafe_b64encode(h.digest()).decode("ascii").rstrip("=") def register_script( self, @@ -964,6 +898,7 @@ def register_script( options: typing.Optional[Options] = None, source_path: typing.Optional[str] = None, module_name: typing.Optional[str] = None, + envs: typing.Optional[typing.Dict[str, str]] = None, ) -> typing.Union[FlyteWorkflow, FlyteTask]: """ Use this method to register a workflow via script mode. @@ -978,6 +913,7 @@ def register_script( :param options: Additional execution options that can be configured for the default launchplan :param source_path: The root of the project path :param module_name: the name of the module + :param envs: Environment variables to be passed to the serialization :return: """ if image_config is None: @@ -998,6 +934,7 @@ def register_script( domain=domain, image_config=image_config, git_repo=_get_git_repo_url(source_path), + env=envs, fast_serialization_settings=FastSerializationSettings( enabled=True, destination_dir=destination_dir, @@ -1109,22 +1046,8 @@ def _execute( ) if isinstance(v, Literal): lit = v - elif isinstance(v, artifacts_pb2.Artifact): - lit = v.spec.value elif isinstance(v, Artifact): - if v.literal is not None: - lit = v.literal - elif v.artifact_id is not None: - fetched_artifact = self.get_artifact(artifact_id=v.artifact_id) - if not fetched_artifact: - raise user_exceptions.FlyteValueException( - v.artifact_id, "Could not find artifact with ID given" - ) - lit = fetched_artifact.literal - else: - raise user_exceptions.FlyteValueException( - v, "When binding input to Artifact, either the Literal or the ID must be set" - ) + raise user_exceptions.FlyteValueException(v, "Running with an artifact object is not yet possible.") else: if k not in type_hints: try: diff --git a/flytekit/sensor/sensor_engine.py b/flytekit/sensor/sensor_engine.py index 79d2e0f4b4..5fa16162a0 100644 --- a/flytekit/sensor/sensor_engine.py +++ b/flytekit/sensor/sensor_engine.py @@ -25,6 +25,8 @@ class SensorEngine(AgentBase): + name = "Sensor" + def __init__(self): super().__init__(task_type="sensor", asynchronous=True) diff --git a/flytekit/tools/fast_registration.py b/flytekit/tools/fast_registration.py index 0664ebbd8d..4b018ce94b 100644 --- a/flytekit/tools/fast_registration.py +++ b/flytekit/tools/fast_registration.py @@ -7,6 +7,7 @@ import subprocess as _subprocess import tarfile import tempfile +import typing from typing import Optional import click @@ -42,7 +43,13 @@ def fast_package(source: os.PathLike, output_dir: os.PathLike, deref_symlinks: b with tempfile.TemporaryDirectory() as tmp_dir: tar_path = os.path.join(tmp_dir, "tmp.tar") with tarfile.open(tar_path, "w", dereference=deref_symlinks) as tar: - tar.add(source, arcname="", filter=lambda x: ignore.tar_filter(tar_strip_file_attributes(x))) + files: typing.List[str] = os.listdir(source) + for ws_file in files: + tar.add( + os.path.join(source, ws_file), + arcname=ws_file, + filter=lambda x: ignore.tar_filter(tar_strip_file_attributes(x)), + ) with gzip.GzipFile(filename=archive_fname, mode="wb", mtime=0) as gzipped: with open(tar_path, "rb") as tar_file: gzipped.write(tar_file.read()) diff --git a/flytekit/tools/repo.py b/flytekit/tools/repo.py index 195e3eea17..c3a456c20b 100644 --- a/flytekit/tools/repo.py +++ b/flytekit/tools/repo.py @@ -97,7 +97,9 @@ def package( click.secho(f"Fast mode enabled: compressed archive {archive_fname}", dim=True) with tarfile.open(output, "w:gz") as tar: - tar.add(output_tmpdir, arcname="") + files: typing.List[str] = os.listdir(output_tmpdir) + for ws_file in files: + tar.add(os.path.join(output_tmpdir, ws_file), arcname=ws_file) click.secho(f"Successfully packaged {len(serializable_entities)} flyte objects into {output}", fg="green") @@ -220,6 +222,7 @@ def register( env: typing.Optional[typing.Dict[str, str]], dry_run: bool = False, activate_launchplans: bool = False, + skip_errors: bool = False, ): detected_root = find_common_root(package_or_module) click.secho(f"Detected Root {detected_root}, using this to create deployable package...", fg="yellow") @@ -276,14 +279,19 @@ def register( secho(og_id, "") try: if not dry_run: - i = remote.raw_register( - cp_entity, serialization_settings, version=version, create_default_launchplan=False - ) - secho(i) - if is_lp and activate_launchplans: - secho(og_id, "", op="Activation") - remote.activate_launchplan(i) - secho(i, reason="activated", op="Activation") + try: + i = remote.raw_register( + cp_entity, serialization_settings, version=version, create_default_launchplan=False + ) + secho(i, state="success") + if is_lp and activate_launchplans: + secho(og_id, "", op="Activation") + remote.activate_launchplan(i) + secho(i, reason="activated", op="Activation") + except Exception as e: + if not skip_errors: + raise e + secho(og_id, state="failed") else: secho(og_id, reason="Dry run Mode!") except RegistrationSkipped: diff --git a/flytekit/tools/script_mode.py b/flytekit/tools/script_mode.py index 5c9990d819..fba454ce76 100644 --- a/flytekit/tools/script_mode.py +++ b/flytekit/tools/script_mode.py @@ -48,7 +48,10 @@ def compress_scripts(source_path: str, destination: str, module_name: str): copy_module_to_destination(source_path, destination_path, module_name, visited) tar_path = os.path.join(tmp_dir, "tmp.tar") with tarfile.open(tar_path, "w") as tar: - tar.add(os.path.join(tmp_dir, "code"), arcname="", filter=tar_strip_file_attributes) + tmp_path: str = os.path.join(tmp_dir, "code") + files: typing.List[str] = os.listdir(tmp_path) + for ws_file in files: + tar.add(os.path.join(tmp_path, ws_file), arcname=ws_file, filter=tar_strip_file_attributes) with gzip.GzipFile(filename=destination, mode="wb", mtime=0) as gzipped: with open(tar_path, "rb") as tar_file: gzipped.write(tar_file.read()) diff --git a/flytekit/tools/translator.py b/flytekit/tools/translator.py index dc81ba55e6..1c2016b681 100644 --- a/flytekit/tools/translator.py +++ b/flytekit/tools/translator.py @@ -159,8 +159,10 @@ def fn(settings: SerializationSettings) -> List[str]: def get_serializable_task( + entity_mapping: OrderedDict, settings: SerializationSettings, entity: FlyteLocalEntity, + options: Optional[Options] = None, ) -> TaskSpec: task_id = _identifier_model.Identifier( _identifier_model.ResourceType.TASK, @@ -176,6 +178,10 @@ def get_serializable_task( # during dynamic serialization settings = settings.with_serialized_context() + if entity.node_dependency_hints is not None: + for entity_hint in entity.node_dependency_hints: + get_serializable(entity_mapping, settings, entity_hint, options) + container = entity.get_container(settings) # This pod will be incorrect when doing fast serialize pod = entity.get_k8s_pod(settings) @@ -714,7 +720,7 @@ def get_serializable( cp_entity = get_reference_spec(entity_mapping, settings, entity) elif isinstance(entity, PythonTask): - cp_entity = get_serializable_task(settings, entity) + cp_entity = get_serializable_task(entity_mapping, settings, entity) elif isinstance(entity, WorkflowBase): cp_entity = get_serializable_workflow(entity_mapping, settings, entity, options) diff --git a/flytekit/trigger.py b/flytekit/trigger.py index 62d703f642..e69de29bb2 100644 --- a/flytekit/trigger.py +++ b/flytekit/trigger.py @@ -1,142 +0,0 @@ -from datetime import timedelta -from typing import Any, Dict, List, Optional, Type, Union - -import isodate -from flyteidl.core import artifact_id_pb2 as art_id -from flyteidl.core import identifier_pb2 as idl -from flyteidl.core import interface_pb2 - -from flytekit.core.artifact import Artifact, ArtifactQuery, Partition, TimePartition -from flytekit.core.context_manager import FlyteContextManager -from flytekit.core.launch_plan import LaunchPlan -from flytekit.core.tracker import TrackedInstance -from flytekit.core.type_engine import TypeEngine -from flytekit.core.workflow import WorkflowBase -from flytekit.models.interface import ParameterMap, Variable - - -class Trigger(TrackedInstance): - """ - Trigger( - trigger_on=[dailyArtifact, hourlyArtifact], - inputs={ - "today_upstream": dailyArtifact, # this means: use the matched artifact - "yesterday_upstream": dailyArtifact.query(time_partition=dailyArtifact.time_partition - timedelta(days=1)), - # this means: use the matched hourly artifact - "other_daily_upstream": hourlyArtifact.query(partitions={"region": "LAX"}), - "region": "SEA", # static value that will be passed as input - "other_artifact": UnrelatedArtifact.query(time_partition=dailyArtifact.time_partition - timedelta(days=1)), - "other_artifact_2": UnrelatedArtifact.query(time_partition=hourlyArtifact.time_partition.truncate_to_day()), - "other_artifact_3": UnrelatedArtifact.query(region=hourlyArtifact.time_partition.truncate_to_day()), - }, - ) - """ - - def __init__( - self, trigger_on: List[Artifact], inputs: Optional[Dict[str, Union[Any, Artifact, ArtifactQuery]]] = None - ): - super().__init__() - for t in trigger_on: - if not isinstance(t, Artifact): - raise ValueError("Trigger must be called with a list of artifacts") - self.triggers = trigger_on - self.inputs = inputs or {} # user doesn't need to specify if the launch plan doesn't take inputs - for k, v in self.inputs.items(): - if isinstance(v, ArtifactQuery): - if v.bindings: - for b in v.bindings: - if b not in self.triggers: - raise ValueError( - f"Binding {b} id {id(b)} must reference a triggering artifact {self.triggers}" - ) - - if isinstance(v, Artifact): - if v not in self.triggers: - raise ValueError(f"Binding {v} id {id(v)} must reference a triggering artifact {self.triggers}") - - def get_parameter_map( - self, input_python_interface: Dict[str, Type], input_typed_interface: Dict[str, Variable] - ) -> interface_pb2.ParameterMap: - """ - This is the key function that enables triggers to work. When declaring a trigger, the user specifies an input - map in the form of artifacts, artifact time partitions, and artifact queries (possibly on unrelated artifacts). - When it comes time to create the trigger, we need to convert all of these into a parameter map (because we've - chosen Parameter as the method by which things like artifact queries are passed around). This function does - that, and converts constants to Literals. - """ - ctx = FlyteContextManager.current_context() - pm = {} - for k, v in self.inputs.items(): - var = input_typed_interface[k].to_flyte_idl() - if isinstance(v, Artifact): - aq = v.embed_as_query(self.triggers) - p = interface_pb2.Parameter(var=var, artifact_query=aq) - elif isinstance(v, ArtifactQuery): - p = interface_pb2.Parameter(var=var, artifact_query=v.to_flyte_idl(self.triggers)) - elif isinstance(v, TimePartition): - expr = None - if v.op and v.other and isinstance(v.other, timedelta): - expr = str(v.op) + isodate.duration_isoformat(v.other) - aq = v.reference_artifact.embed_as_query(self.triggers, bind_to_time_partition=True, expr=expr) - p = interface_pb2.Parameter(var=var, artifact_query=aq) - elif isinstance(v, Partition): - # The reason is that if we bind to arbitrary partitions, we'll have to start keeping track of types - # and if the experiment service is written in non-python, we'd have to reimplement a type engine in - # the other language - raise ValueError( - "Don't bind to non-time partitions. Time partitions are okay because of the known type." - ) - else: - lit = TypeEngine.to_literal(ctx, v, input_python_interface[k], input_typed_interface[k].type) - p = interface_pb2.Parameter(var=var, default=lit.to_flyte_idl()) - - pm[k] = p - return interface_pb2.ParameterMap(parameters=pm) - - def to_flyte_idl(self) -> art_id.Trigger: - try: - name = f"{self.instantiated_in}.{self.lhs}" - except Exception: # noqa broad for now given the changing nature of the tracker implementation. - import random - from uuid import UUID - - name = "trigger" + UUID(int=random.getrandbits(128)).hex - - # project/domain will be empty - to be bound later at registration time. - artifact_ids = [a.to_flyte_idl().artifact_id for a in self.triggers] - - return art_id.Trigger( - trigger_id=idl.Identifier( - resource_type=idl.ResourceType.LAUNCH_PLAN, - name=name, - ), - triggers=artifact_ids, - ) - - def __call__(self, *args, **kwargs): - if len(kwargs) > 0: - raise ValueError("Trigger does not support keyword arguments") - if len(args) != 1: - raise ValueError("Trigger must be called with a single function") - - entity = args[0] - # This needs to create a Launch plan in the same - if not isinstance(entity, WorkflowBase) and not isinstance(entity, LaunchPlan): - raise ValueError("Trigger must be called with a workflow or launch plan") - - ctx = FlyteContextManager.current_context() - - pm = self.get_parameter_map(entity.python_interface.inputs, entity.interface.inputs) - - pm_model = ParameterMap.from_flyte_idl(pm) - - if isinstance(entity, LaunchPlan): - raise NotImplementedError("Launch plan triggers are not yet supported") - - if isinstance(entity, WorkflowBase): - default_lp = LaunchPlan.get_default_launch_plan(ctx, entity) - trigger_lp = default_lp.clone_with(name=default_lp.name, parameters=pm_model) - self_idl = self.to_flyte_idl() - trigger_lp._additional_metadata = self_idl - - return entity diff --git a/plugins/flytekit-airflow/dev-requirements.txt b/plugins/flytekit-airflow/dev-requirements.txt index b7b8178724..114279f520 100644 --- a/plugins/flytekit-airflow/dev-requirements.txt +++ b/plugins/flytekit-airflow/dev-requirements.txt @@ -6,7 +6,7 @@ # aiofiles==23.2.1 # via gcloud-aio-storage -aiohttp==3.9.0 +aiohttp==3.9.2 # via # apache-airflow-providers-http # gcloud-aio-auth diff --git a/plugins/flytekit-airflow/flytekitplugins/airflow/agent.py b/plugins/flytekit-airflow/flytekitplugins/airflow/agent.py index 02cf70cbae..b84deb009d 100644 --- a/plugins/flytekit-airflow/flytekitplugins/airflow/agent.py +++ b/plugins/flytekit-airflow/flytekitplugins/airflow/agent.py @@ -62,6 +62,8 @@ class AirflowAgent(AgentBase): In this case, those operators will be converted to AirflowContainerTask and executed in the pod. """ + name = "Airflow Agent" + def __init__(self): super().__init__(task_type="airflow", asynchronous=True) diff --git a/plugins/flytekit-bigquery/flytekitplugins/bigquery/agent.py b/plugins/flytekit-bigquery/flytekitplugins/bigquery/agent.py index baf3a32500..54b70ea8e0 100644 --- a/plugins/flytekit-bigquery/flytekitplugins/bigquery/agent.py +++ b/plugins/flytekit-bigquery/flytekitplugins/bigquery/agent.py @@ -43,6 +43,8 @@ class Metadata: class BigQueryAgent(AgentBase): + name = "Bigquery Agent" + def __init__(self): super().__init__(task_type="bigquery_query_job_task", asynchronous=False) diff --git a/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py b/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py index 45d0b4676c..6c11bd6838 100644 --- a/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py +++ b/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py @@ -1,8 +1,8 @@ -import json import os import pathlib import shutil import subprocess +from importlib import metadata import click from packaging.version import Version @@ -108,9 +108,7 @@ def build(): if image_spec.source_root: shutil.copytree(image_spec.source_root, pathlib.Path(cfg_path).parent, dirs_exist_ok=True) - version_command = "envd version -s -f json" - envd_version = json.loads(EnvdImageSpecBuilder().execute_command(version_command)[0])["envd"].replace("v", "") - + envd_version = metadata.version("envd") # Indentation is required by envd if Version(envd_version) <= Version("0.3.37"): envd_config += ' io.copy(host_path="./", envd_path="/root")' diff --git a/plugins/flytekit-envd/tests/test_image_spec.py b/plugins/flytekit-envd/tests/test_image_spec.py index 36adebd346..d77b2ca89b 100644 --- a/plugins/flytekit-envd/tests/test_image_spec.py +++ b/plugins/flytekit-envd/tests/test_image_spec.py @@ -1,9 +1,23 @@ from pathlib import Path from textwrap import dedent +import pytest from flytekitplugins.envd.image_builder import EnvdImageSpecBuilder, create_envd_config -from flytekit.image_spec.image_spec import ImageSpec +from flytekit.image_spec.image_spec import ImageBuildEngine, ImageSpec + + +@pytest.fixture(scope="module", autouse=True) +def register_envd_higher_priority(): + # Register a new envd platform with the highest priority so the test in this file uses envd + highest_priority_builder = max(ImageBuildEngine._REGISTRY, key=ImageBuildEngine._REGISTRY.get) + highest_priority = ImageBuildEngine._REGISTRY[highest_priority_builder][1] + yield ImageBuildEngine.register( + "envd_high_priority", + EnvdImageSpecBuilder(), + priority=highest_priority + 1, + ) + del ImageBuildEngine._REGISTRY["envd_high_priority"] def test_image_spec(): diff --git a/plugins/flytekit-flyin/tests/test_flyin_plugin.py b/plugins/flytekit-flyin/tests/test_flyin_plugin.py index 23cde516ce..ce2758e595 100644 --- a/plugins/flytekit-flyin/tests/test_flyin_plugin.py +++ b/plugins/flytekit-flyin/tests/test_flyin_plugin.py @@ -1,3 +1,5 @@ +from collections import OrderedDict + import mock import pytest from flytekitplugins.flyin import ( @@ -343,7 +345,7 @@ def t(): project="p", domain="d", version="v", image_config=default_image_config ) - serialized_task = get_serializable_task(default_serialization_settings, t) + serialized_task = get_serializable_task(OrderedDict(), default_serialization_settings, t) assert serialized_task.template.config == {"link_type": "vscode", "port": "8081"} diff --git a/plugins/flytekit-mmcloud/flytekitplugins/mmcloud/agent.py b/plugins/flytekit-mmcloud/flytekitplugins/mmcloud/agent.py index b44906e144..5c6ab06831 100644 --- a/plugins/flytekit-mmcloud/flytekitplugins/mmcloud/agent.py +++ b/plugins/flytekit-mmcloud/flytekitplugins/mmcloud/agent.py @@ -22,8 +22,10 @@ class Metadata: class MMCloudAgent(AgentBase): + name = "MMCloud Agent" + def __init__(self): - super().__init__(task_type="mmcloud_task") + super().__init__(task_type="mmcloud_task", asynchronous=True) self._response_format = ["--format", "json"] async def async_login(self): diff --git a/plugins/flytekit-onnx-pytorch/dev-requirements.txt b/plugins/flytekit-onnx-pytorch/dev-requirements.txt index 44660a3ba9..35a9b49a8f 100644 --- a/plugins/flytekit-onnx-pytorch/dev-requirements.txt +++ b/plugins/flytekit-onnx-pytorch/dev-requirements.txt @@ -69,7 +69,7 @@ onnxruntime==1.16.1 # via -r dev-requirements.in packaging==23.2 # via onnxruntime -pillow==10.1.0 +pillow==10.2.0 # via # -r dev-requirements.in # torchvision diff --git a/plugins/flytekit-onnx-tensorflow/dev-requirements.txt b/plugins/flytekit-onnx-tensorflow/dev-requirements.txt index 689a985a98..38a63b116c 100644 --- a/plugins/flytekit-onnx-tensorflow/dev-requirements.txt +++ b/plugins/flytekit-onnx-tensorflow/dev-requirements.txt @@ -18,7 +18,7 @@ onnxruntime==1.16.1 # via -r dev-requirements.in packaging==23.2 # via onnxruntime -pillow==10.1.0 +pillow==10.2.0 # via -r dev-requirements.in protobuf==4.25.0 # via onnxruntime diff --git a/plugins/flytekit-papermill/dev-requirements.in b/plugins/flytekit-papermill/dev-requirements.in index 88425c012f..d0a9617bdb 100644 --- a/plugins/flytekit-papermill/dev-requirements.in +++ b/plugins/flytekit-papermill/dev-requirements.in @@ -1,4 +1,4 @@ -flyteidl>=1.10.0 +flyteidl>=1.10.7b0 -e file:../../.#egg=flytekitplugins-pod&subdirectory=plugins/flytekit-k8s-pod -e file:../../.#egg=flytekitplugins-spark&subdirectory=plugins/flytekit-spark -e file:../../.#egg=flytekitplugins-awsbatch&subdirectory=plugins/flytekit-aws-batch diff --git a/plugins/flytekit-papermill/dev-requirements.txt b/plugins/flytekit-papermill/dev-requirements.txt index 5928099f27..21c03fd02d 100644 --- a/plugins/flytekit-papermill/dev-requirements.txt +++ b/plugins/flytekit-papermill/dev-requirements.txt @@ -93,7 +93,7 @@ docker-image-py==0.1.12 # via flytekit docstring-parser==0.15 # via flytekit -flyteidl==1.10.6 +flyteidl==1.10.7b0 # via # -r dev-requirements.in # flytekit @@ -235,7 +235,9 @@ packaging==23.2 # docker # marshmallow pandas==1.5.3 - # via flytekit + # via + # flytekit + # flytekitplugins-spark portalocker==2.8.2 # via msal-extensions protobuf==4.24.4 diff --git a/plugins/flytekit-spark/dev-requirements.txt b/plugins/flytekit-spark/dev-requirements.txt index 3335091569..5f5f8e283a 100644 --- a/plugins/flytekit-spark/dev-requirements.txt +++ b/plugins/flytekit-spark/dev-requirements.txt @@ -4,7 +4,7 @@ # # pip-compile dev-requirements.in # -aiohttp==3.9.1 +aiohttp==3.9.2 # via aioresponses aioresponses==0.7.6 # via -r dev-requirements.in diff --git a/plugins/flytekit-spark/flytekitplugins/spark/agent.py b/plugins/flytekit-spark/flytekitplugins/spark/agent.py index ffaa11ee13..ff907d2a41 100644 --- a/plugins/flytekit-spark/flytekitplugins/spark/agent.py +++ b/plugins/flytekit-spark/flytekitplugins/spark/agent.py @@ -26,8 +26,10 @@ class Metadata: class DatabricksAgent(AgentBase): + name = "Databricks Agent" + def __init__(self): - super().__init__(task_type="spark") + super().__init__(task_type="spark", asynchronous=True) async def async_create( self, diff --git a/pyproject.toml b/pyproject.toml index 9ca9119d64..1c57b86cb8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,7 +21,7 @@ dependencies = [ "diskcache>=5.2.1", "docker>=4.0.0,<7.0.0", "docstring-parser>=0.9.0", - "flyteidl>=1.10.6", + "flyteidl>=1.10.7b2", "fsspec>=2023.3.0,<=2023.9.2", "gcsfs>=2023.3.0,<=2023.9.2", "googleapis-common-protos>=1.57", @@ -36,8 +36,7 @@ dependencies = [ "marshmallow-enum", "marshmallow-jsonschema>=0.12.0", "mashumaro>=3.9.1", - # TODO: Remove upper-bound after protobuf community fixes it. https://github.com/flyteorg/flyte/issues/4359 - "protobuf<4.25.0", + "protobuf!=4.25.0", "pyarrow", "python-json-logger>=2.0.0", "pytimeparse>=1.1.8,<2.0.0", @@ -58,6 +57,7 @@ classifiers = [ "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", "Topic :: Scientific/Engineering", "Topic :: Scientific/Engineering :: Artificial Intelligence", "Topic :: Software Development", diff --git a/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt b/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt index 5f4ef273a3..571db39e05 100644 --- a/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt +++ b/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt @@ -8,7 +8,7 @@ adlfs==2023.4.0 # via flytekit aiobotocore==2.5.2 # via s3fs -aiohttp==3.9.1 +aiohttp==3.9.2 # via # adlfs # aiobotocore diff --git a/tests/flytekit/unit/cli/pyflyte/test_run.py b/tests/flytekit/unit/cli/pyflyte/test_run.py index bab8b95b15..b42be184f8 100644 --- a/tests/flytekit/unit/cli/pyflyte/test_run.py +++ b/tests/flytekit/unit/cli/pyflyte/test_run.py @@ -198,14 +198,24 @@ def test_union_type_with_invalid_input(): def test_get_entities_in_file(): e = get_entities_in_file(WORKFLOW_FILE, False) - assert e.workflows == ["my_wf", "wf_with_none"] - assert e.tasks == ["get_subset_df", "print_all", "show_sd", "task_with_optional", "test_union1", "test_union2"] + assert e.workflows == ["my_wf", "wf_with_env_vars", "wf_with_none"] + assert e.tasks == [ + "get_subset_df", + "print_all", + "show_sd", + "task_with_env_vars", + "task_with_optional", + "test_union1", + "test_union2", + ] assert e.all() == [ "my_wf", + "wf_with_env_vars", "wf_with_none", "get_subset_df", "print_all", "show_sd", + "task_with_env_vars", "task_with_optional", "test_union1", "test_union2", @@ -390,3 +400,33 @@ def test_pyflyte_run_with_none(a_val): else: assert output == a_val assert result.exit_code == 0 + + +@pytest.mark.parametrize( + "envs, envs_argument, expected_output", + [ + (["--env", "MY_ENV_VAR=hello"], '["MY_ENV_VAR"]', "hello"), + (["--env", "MY_ENV_VAR=hello", "--env", "ABC=42"], '["MY_ENV_VAR","ABC"]', "hello,42"), + ], +) +def test_envvar_local_execution(envs, envs_argument, expected_output): + runner = CliRunner() + args = ( + [ + "run", + ] + + envs + + [ + WORKFLOW_FILE, + "wf_with_env_vars", + "--env_vars", + ] + + [envs_argument] + ) + result = runner.invoke( + pyflyte.main, + args, + catch_exceptions=False, + ) + output = result.stdout.strip().split("\n")[-1].strip() + assert output == expected_output diff --git a/tests/flytekit/unit/cli/pyflyte/workflow.py b/tests/flytekit/unit/cli/pyflyte/workflow.py index 59b0e1b4b2..95535d2fc0 100644 --- a/tests/flytekit/unit/cli/pyflyte/workflow.py +++ b/tests/flytekit/unit/cli/pyflyte/workflow.py @@ -1,5 +1,6 @@ import datetime import enum +import os import typing from dataclasses import dataclass @@ -111,3 +112,16 @@ def task_with_optional(a: typing.Optional[str]) -> str: @workflow def wf_with_none(a: typing.Optional[str] = None) -> str: return task_with_optional(a=a) + + +@task +def task_with_env_vars(env_vars: typing.List[str]) -> str: + collated_env_vars = [] + for env_var in env_vars: + collated_env_vars.append(os.environ[env_var]) + return ",".join(collated_env_vars) + + +@workflow +def wf_with_env_vars(env_vars: typing.List[str]) -> str: + return task_with_env_vars(env_vars=env_vars) diff --git a/tests/flytekit/unit/configuration/test_image_config.py b/tests/flytekit/unit/configuration/test_image_config.py index 7e3c8ad300..82aea000b0 100644 --- a/tests/flytekit/unit/configuration/test_image_config.py +++ b/tests/flytekit/unit/configuration/test_image_config.py @@ -1,9 +1,11 @@ import os import sys +from unittest.mock import Mock import mock import pytest +import flytekit from flytekit.configuration import ImageConfig from flytekit.configuration.default_images import DefaultImages, PythonVersion @@ -63,3 +65,14 @@ def test_image_create(): def test_get_version_suffix(): assert DefaultImages.get_version_suffix() == "latest" + + +def test_default_image_plugin(monkeypatch): + new_default_image = "registry/flytekit:py3.9-latest" + + plugin_mock = Mock() + plugin_mock.get_default_image.return_value = new_default_image + mock_global_plugin = {"plugin": plugin_mock} + monkeypatch.setattr(flytekit.configuration.plugin, "_GLOBAL_CONFIG", mock_global_plugin) + + assert DefaultImages.default_image() == new_default_image diff --git a/tests/flytekit/unit/core/image_spec/test_image_spec.py b/tests/flytekit/unit/core/image_spec/test_image_spec.py index 81151beddd..38727d02fd 100644 --- a/tests/flytekit/unit/core/image_spec/test_image_spec.py +++ b/tests/flytekit/unit/core/image_spec/test_image_spec.py @@ -1,4 +1,5 @@ import os +from unittest.mock import Mock import pytest @@ -75,3 +76,24 @@ def test_image_spec(mock_image_spec_builder): # ImageSpec should be immutable image_spec.with_commands("ls") assert image_spec.commands == ["echo hello"] + + +def test_image_spec_engine_priority(): + image_spec = ImageSpec(name="FLYTEKIT") + image_name = image_spec.image_name() + + new_image_name = f"fqn.xyz/{image_name}" + mock_image_builder_10 = Mock() + mock_image_builder_10.build_image.return_value = new_image_name + mock_image_builder_default = Mock() + mock_image_builder_default.build_image.side_effect = ValueError("should not be called") + + ImageBuildEngine.register("build_10", mock_image_builder_10, priority=10) + ImageBuildEngine.register("build_default", mock_image_builder_default) + + ImageBuildEngine.build(image_spec) + mock_image_builder_10.build_image.assert_called_once_with(image_spec) + + assert image_spec.image_name() == new_image_name + del ImageBuildEngine._REGISTRY["build_10"] + del ImageBuildEngine._REGISTRY["build_default"] diff --git a/tests/flytekit/unit/core/test_array_node_map_task.py b/tests/flytekit/unit/core/test_array_node_map_task.py index aeb727f5f8..40bb864c4f 100644 --- a/tests/flytekit/unit/core/test_array_node_map_task.py +++ b/tests/flytekit/unit/core/test_array_node_map_task.py @@ -7,7 +7,7 @@ from flytekit import task, workflow from flytekit.configuration import FastSerializationSettings, Image, ImageConfig, SerializationSettings -from flytekit.core.array_node_map_task import ArrayNodeMapTask +from flytekit.core.array_node_map_task import ArrayNodeMapTask, ArrayNodeMapTaskResolver from flytekit.core.task import TaskMetadata from flytekit.experimental import map_task as array_node_map_task from flytekit.tools.translator import get_serializable @@ -187,7 +187,7 @@ def many_inputs(a: int, b: str, c: float) -> str: assert m.python_interface.inputs == {"a": List[int], "b": List[str], "c": List[float]} assert ( m.name - == "tests.flytekit.unit.core.test_array_node_map_task.map_many_inputs_4ee240ef5cf979dbc133fb30035cb874-arraynode" + == "tests.flytekit.unit.core.test_array_node_map_task.map_many_inputs_bf51001578d0ae197a52c0af0a99dd89-arraynode" ) r_m = ArrayNodeMapTask(many_inputs) assert str(r_m.python_interface) == str(m.python_interface) @@ -197,7 +197,7 @@ def many_inputs(a: int, b: str, c: float) -> str: assert m.python_interface.inputs == {"a": List[int], "b": List[str], "c": float} assert ( m.name - == "tests.flytekit.unit.core.test_array_node_map_task.map_many_inputs_352fcdea8523a83134b51bbf5793f14e-arraynode" + == "tests.flytekit.unit.core.test_array_node_map_task.map_many_inputs_cb470e880fabd6265ec80e29fe60250d-arraynode" ) r_m = ArrayNodeMapTask(many_inputs, bound_inputs=set("c")) assert str(r_m.python_interface) == str(m.python_interface) @@ -207,7 +207,7 @@ def many_inputs(a: int, b: str, c: float) -> str: assert m.python_interface.inputs == {"a": List[int], "b": str, "c": float} assert ( m.name - == "tests.flytekit.unit.core.test_array_node_map_task.map_many_inputs_e224ba3a5b00e08083d541a6ca99b179-arraynode" + == "tests.flytekit.unit.core.test_array_node_map_task.map_many_inputs_316e10eb97f5d2abd585951048b807b9-arraynode" ) r_m = ArrayNodeMapTask(many_inputs, bound_inputs={"c", "b"}) assert str(r_m.python_interface) == str(m.python_interface) @@ -217,7 +217,7 @@ def many_inputs(a: int, b: str, c: float) -> str: assert m.python_interface.inputs == {"a": int, "b": str, "c": float} assert ( m.name - == "tests.flytekit.unit.core.test_array_node_map_task.map_many_inputs_f080e60be9d6faedeef0c74834d6812a-arraynode" + == "tests.flytekit.unit.core.test_array_node_map_task.map_many_inputs_758022acd59ad1c8b81670378d4de4f6-arraynode" ) r_m = ArrayNodeMapTask(many_inputs, bound_inputs={"a", "c", "b"}) assert str(r_m.python_interface) == str(m.python_interface) @@ -257,6 +257,18 @@ def task3(c: str, a: int, b: float) -> str: assert m1 == m2 == m3 == ["1 - 0.1 - c", "2 - 0.2 - c", "3 - 0.3 - c"] +def test_bounded_inputs_vars_order(serialization_settings): + @task() + def task1(a: int, b: float, c: str) -> str: + return f"{a} - {b} - {c}" + + mt = array_node_map_task(functools.partial(task1, c=1.0, b="hello", a=1)) + mtr = ArrayNodeMapTaskResolver() + args = mtr.loader_args(serialization_settings, mt) + + assert args[1] == "a,b,c" + + @pytest.mark.parametrize( "min_success_ratio, should_raise_error", [ diff --git a/tests/flytekit/unit/core/test_artifacts.py b/tests/flytekit/unit/core/test_artifacts.py index dbafa6983f..45a7d07e0a 100644 --- a/tests/flytekit/unit/core/test_artifacts.py +++ b/tests/flytekit/unit/core/test_artifacts.py @@ -1,19 +1,24 @@ import datetime +import sys from collections import OrderedDict -import pandas as pd import pytest from flyteidl.core import artifact_id_pb2 as art_id -from typing_extensions import Annotated +from typing_extensions import Annotated, get_args from flytekit.configuration import Image, ImageConfig, SerializationSettings from flytekit.core.artifact import Artifact, Inputs from flytekit.core.context_manager import FlyteContextManager +from flytekit.core.interface import detect_artifact from flytekit.core.launch_plan import LaunchPlan from flytekit.core.task import task from flytekit.core.workflow import workflow from flytekit.tools.translator import get_serializable +if "pandas" not in sys.modules: + pytest.skip(msg="Requires pandas", allow_module_level=True) + + default_img = Image(name="default", fqn="test", tag="tag") serialization_settings = SerializationSettings( project="project", @@ -30,6 +35,8 @@ def __init__(self, data): def test_basic_option_a_rev(): + import pandas as pd + a1_t_ab = Artifact(name="my_data", partition_keys=["a", "b"], time_partitioned=True) @task @@ -54,7 +61,22 @@ def t1( assert t1_s.template.interface.outputs["o0"].artifact_partial_id.artifact_key.project == "" +def test_args_getting(): + a1 = Artifact(name="argstst") + a1_called = a1() + x = Annotated[int, a1_called] + gotten = get_args(x) + assert len(gotten) == 2 + assert gotten[1] is a1_called + detected = detect_artifact(get_args(int)) + assert detected is None + detected = detect_artifact(get_args(x)) + assert detected == a1_called.to_partial_artifact_id() + + def test_basic_option_no_tp(): + import pandas as pd + a1_t_ab = Artifact(name="my_data", partition_keys=["a", "b"]) assert not a1_t_ab.time_partitioned @@ -92,6 +114,8 @@ def test_basic_option_hardcoded_tp(): def test_basic_option_a(): + import pandas as pd + a1_t_ab = Artifact(name="my_data", partition_keys=["a", "b"], time_partitioned=True) @task @@ -111,6 +135,8 @@ def t1( def test_basic_no_call(): + import pandas as pd + a1_t_ab = Artifact(name="my_data", partition_keys=["a", "b"], time_partitioned=True) # raise an error because the user hasn't () the artifact @@ -123,6 +149,8 @@ def t1(b_value: str, dt: datetime.datetime) -> Annotated[pd.DataFrame, a1_t_ab]: def test_basic_option_a2(): + import pandas as pd + a2_ab = Artifact(name="my_data2", partition_keys=["a", "b"]) with pytest.raises(ValueError): @@ -144,6 +172,8 @@ def t2(b_value: str) -> Annotated[pd.DataFrame, a2_ab(a=Inputs.b_value, b="manua def test_basic_option_a3(): + import pandas as pd + a3 = Artifact(name="my_data3") @task @@ -181,7 +211,7 @@ def test_not_specified_behavior(): assert aq.artifact_id.artifact_key.project == "pr" assert aq.artifact_id.artifact_key.domain == "dom" - assert wf_artifact_no_tag.as_artifact_id.HasField("partitions") is False + assert wf_artifact_no_tag.concrete_artifact_id.HasField("partitions") is False wf_artifact_no_tag = Artifact(project="project1", domain="dev", name="wf_artifact", partitions={}) assert wf_artifact_no_tag.partitions is None diff --git a/tests/flytekit/unit/core/test_artifacts_sandbox.py b/tests/flytekit/unit/core/test_artifacts_sandbox.py index a93dd34d33..648f428260 100644 --- a/tests/flytekit/unit/core/test_artifacts_sandbox.py +++ b/tests/flytekit/unit/core/test_artifacts_sandbox.py @@ -1,4 +1,3 @@ -import pandas as pd import pytest from flyteidl.artifact import artifacts_pb2 from flyteidl.core import identifier_pb2 @@ -39,6 +38,8 @@ def test_create_an_artifact33_locally(): @pytest.mark.sandbox_test def test_create_an_artifact_locally(): + import pandas as pd + df = pd.DataFrame({"Name": ["Mary", "Jane"], "Age": [22, 23]}) # a = Artifact.initialize(python_val=df, python_type=pd.DataFrame, name="flyteorg.test.yt.test1", # aliases=["v0.1.0"]) diff --git a/tests/flytekit/unit/core/test_container_task.py b/tests/flytekit/unit/core/test_container_task.py index 3ac4a47cc4..c89ec11345 100644 --- a/tests/flytekit/unit/core/test_container_task.py +++ b/tests/flytekit/unit/core/test_container_task.py @@ -1,3 +1,5 @@ +from collections import OrderedDict + import pytest from kubernetes.client.models import ( V1Affinity, @@ -72,7 +74,7 @@ def test_pod_template(): ################# # Test Serialization ################# - ts = get_serializable_task(default_serialization_settings, ct) + ts = get_serializable_task(OrderedDict(), default_serialization_settings, ct) assert ts.template.metadata.pod_template_name == "my-base-template" assert ts.template.container is None assert ts.template.k8s_pod is not None diff --git a/tests/flytekit/unit/core/test_dynamic.py b/tests/flytekit/unit/core/test_dynamic.py index b9b0ebd3fa..44cc19b5c3 100644 --- a/tests/flytekit/unit/core/test_dynamic.py +++ b/tests/flytekit/unit/core/test_dynamic.py @@ -1,4 +1,5 @@ import typing +from collections import OrderedDict import pytest @@ -13,6 +14,7 @@ from flytekit.core.type_engine import TypeEngine from flytekit.core.workflow import workflow from flytekit.models.literals import LiteralMap +from flytekit.tools.translator import get_serializable_task settings = flytekit.configuration.SerializationSettings( project="test_proj", @@ -262,3 +264,29 @@ def wf(wf_in: str) -> typing.List[str]: res = dt(ss="hello") assert res == ["In t2 string is hello", "In t3 string is In t2 string is hello"] + + +def test_node_dependency_hints_are_serialized(): + @task + def t1() -> int: + return 0 + + @task + def t2() -> int: + return 0 + + @dynamic(node_dependency_hints=[t1, t2]) + def dt(mode: int) -> int: + if mode == 1: + return t1() + if mode == 2: + return t2() + + raise ValueError("Invalid mode") + + entity_mapping = OrderedDict() + get_serializable_task(entity_mapping, settings, dt) + + serialised_entities_iterator = iter(entity_mapping.values()) + assert "t1" in next(serialised_entities_iterator).template.id.name + assert "t2" in next(serialised_entities_iterator).template.id.name diff --git a/tests/flytekit/unit/core/test_local_cache.py b/tests/flytekit/unit/core/test_local_cache.py index c79930148f..53c524e95f 100644 --- a/tests/flytekit/unit/core/test_local_cache.py +++ b/tests/flytekit/unit/core/test_local_cache.py @@ -100,6 +100,25 @@ def check_evenness(n: int) -> bool: assert n_cached_task_calls == 2 +def test_cache_can_be_disabled(monkeypatch): + monkeypatch.setenv("FLYTE_LOCAL_CACHE_ENABLED", "false") + + @task(cache=True, cache_version="v1") + def is_even(n: int) -> bool: + global n_cached_task_calls + n_cached_task_calls += 1 + return n % 2 == 0 + + assert n_cached_task_calls == 0 + # Run once and check that the counter is increased + assert is_even(n=1) is False + assert n_cached_task_calls == 1 + + # Run again and check that the counter is increased again i.e. no caching + assert is_even(n=1) is False + assert n_cached_task_calls == 2 + + def test_shared_tasks_in_two_separate_workflows(): @task(cache=True, cache_version="0.0.1") def is_odd(n: int) -> bool: diff --git a/tests/flytekit/unit/core/test_map_task.py b/tests/flytekit/unit/core/test_map_task.py index 26d1a71c3c..c87d4c6b1f 100644 --- a/tests/flytekit/unit/core/test_map_task.py +++ b/tests/flytekit/unit/core/test_map_task.py @@ -192,31 +192,36 @@ def many_inputs(a: int, b: str, c: float) -> str: m = map_task(many_inputs) assert m.python_interface.inputs == {"a": typing.List[int], "b": typing.List[str], "c": typing.List[float]} - assert m.name == "tests.flytekit.unit.core.test_map_task.map_many_inputs_24c08b3a2f9c2e389ad9fc6a03482cf9" + assert m.name == "tests.flytekit.unit.core.test_map_task.map_many_inputs_d41d8cd98f00b204e9800998ecf8427e" r_m = MapPythonTask(many_inputs) assert str(r_m.python_interface) == str(m.python_interface) p1 = functools.partial(many_inputs, c=1.0) m = map_task(p1) assert m.python_interface.inputs == {"a": typing.List[int], "b": typing.List[str], "c": float} - assert m.name == "tests.flytekit.unit.core.test_map_task.map_many_inputs_697aa7389996041183cf6cfd102be4f7" + assert m.name == "tests.flytekit.unit.core.test_map_task.map_many_inputs_4a8a08f09d37b73795649038408b5f33" r_m = MapPythonTask(many_inputs, bound_inputs=set("c")) assert str(r_m.python_interface) == str(m.python_interface) p2 = functools.partial(p1, b="hello") m = map_task(p2) assert m.python_interface.inputs == {"a": typing.List[int], "b": str, "c": float} - assert m.name == "tests.flytekit.unit.core.test_map_task.map_many_inputs_cc18607da7494024a402a5fa4b3ea5c6" + assert m.name == "tests.flytekit.unit.core.test_map_task.map_many_inputs_74aefa13d6ab8e4bfbd241583749dfe8" r_m = MapPythonTask(many_inputs, bound_inputs={"c", "b"}) assert str(r_m.python_interface) == str(m.python_interface) p3 = functools.partial(p2, a=1) m = map_task(p3) assert m.python_interface.inputs == {"a": int, "b": str, "c": float} - assert m.name == "tests.flytekit.unit.core.test_map_task.map_many_inputs_52fe80b04781ea77ef6f025f4b49abef" + assert m.name == "tests.flytekit.unit.core.test_map_task.map_many_inputs_a44c56c8177e32d3613988f4dba7962e" r_m = MapPythonTask(many_inputs, bound_inputs={"a", "c", "b"}) assert str(r_m.python_interface) == str(m.python_interface) + p3_1 = functools.partial(p2, a=1) + m_1 = map_task(p3_1) + assert m_1.python_interface.inputs == {"a": int, "b": str, "c": float} + assert m_1.name == m.name + with pytest.raises(TypeError): m(a=[1, 2, 3]) @@ -348,3 +353,11 @@ def wf(x: typing.List[int]): map_task(my_mappable_task)(a=x).with_overrides(container_image="random:image") assert wf.nodes[0].flyte_entity.run_task.container_image == "random:image" + + +def test_bounded_inputs_vars_order(serialization_settings): + mt = map_task(functools.partial(t3, c=1.0, b="hello", a=1)) + mtr = MapTaskResolver() + args = mtr.loader_args(serialization_settings, mt) + + assert args[1] == "a,b,c" diff --git a/tests/flytekit/unit/core/test_python_auto_container.py b/tests/flytekit/unit/core/test_python_auto_container.py index fed612c98c..f5c7136b0b 100644 --- a/tests/flytekit/unit/core/test_python_auto_container.py +++ b/tests/flytekit/unit/core/test_python_auto_container.py @@ -1,3 +1,4 @@ +from collections import OrderedDict from typing import Any import pytest @@ -73,7 +74,7 @@ def test_get_container(default_serialization_settings): assert c.image == "docker.io/xyz:some-git-hash" assert c.env == {"FOO": "bar"} - ts = get_serializable_task(default_serialization_settings, task) + ts = get_serializable_task(OrderedDict(), default_serialization_settings, task) assert ts.template.container.image == "docker.io/xyz:some-git-hash" assert ts.template.container.env == {"FOO": "bar"} @@ -86,7 +87,7 @@ def test_get_container_with_task_envvars(default_serialization_settings): assert c.image == "docker.io/xyz:some-git-hash" assert c.env == {"FOO": "bar", "HAM": "spam"} - ts = get_serializable_task(default_serialization_settings, task_with_env_vars) + ts = get_serializable_task(OrderedDict(), default_serialization_settings, task_with_env_vars) assert ts.template.container.image == "docker.io/xyz:some-git-hash" assert ts.template.container.env == {"FOO": "bar", "HAM": "spam"} @@ -96,7 +97,7 @@ def test_get_container_without_serialization_settings_envvars(minimal_serializat assert c.image == "docker.io/xyz:some-git-hash" assert c.env == {"HAM": "spam"} - ts = get_serializable_task(minimal_serialization_settings, task_with_env_vars) + ts = get_serializable_task(OrderedDict(), minimal_serialization_settings, task_with_env_vars) assert ts.template.container.image == "docker.io/xyz:some-git-hash" assert ts.template.container.env == {"HAM": "spam"} @@ -215,7 +216,7 @@ def test_pod_template(default_serialization_settings): ################# # Test Serialization ################# - ts = get_serializable_task(default_serialization_settings, task_with_pod_template) + ts = get_serializable_task(OrderedDict(), default_serialization_settings, task_with_pod_template) assert ts.template.container is None # k8s_pod content is already verified above, so only check the existence here assert ts.template.k8s_pod is not None @@ -290,7 +291,7 @@ def test_minimum_pod_template(default_serialization_settings): ################# # Test Serialization ################# - ts = get_serializable_task(default_serialization_settings, task_with_minimum_pod_template) + ts = get_serializable_task(OrderedDict(), default_serialization_settings, task_with_minimum_pod_template) assert ts.template.container is None # k8s_pod content is already verified above, so only check the existence here assert ts.template.k8s_pod is not None diff --git a/tests/flytekit/unit/core/test_python_function_task.py b/tests/flytekit/unit/core/test_python_function_task.py index 498228f3fe..bec160cd7e 100644 --- a/tests/flytekit/unit/core/test_python_function_task.py +++ b/tests/flytekit/unit/core/test_python_function_task.py @@ -1,3 +1,5 @@ +from collections import OrderedDict + import pytest from kubernetes.client.models import V1Container, V1PodSpec @@ -205,8 +207,20 @@ def func_with_pod_template(i: str): ################# # Test Serialization ################# - ts = get_serializable_task(default_serialization_settings, func_with_pod_template) + ts = get_serializable_task(OrderedDict(), default_serialization_settings, func_with_pod_template) assert ts.template.container is None # k8s_pod content is already verified above, so only check the existence here assert ts.template.k8s_pod is not None assert ts.template.metadata.pod_template_name == "A" + + +def test_node_dependency_hints_are_not_allowed(): + @task + def t1(i: str): + pass + + with pytest.raises(ValueError, match="node_dependency_hints should only be used on dynamic tasks"): + + @task(node_dependency_hints=[t1]) + def t2(i: str): + pass diff --git a/tests/flytekit/unit/core/test_triggers.py b/tests/flytekit/unit/core/test_triggers.py index d00067b7d0..e69de29bb2 100644 --- a/tests/flytekit/unit/core/test_triggers.py +++ b/tests/flytekit/unit/core/test_triggers.py @@ -1,227 +0,0 @@ -from datetime import datetime, timedelta - -from flyteidl.core import artifact_id_pb2 as art_id -from flyteidl.core import literals_pb2 -from typing_extensions import Annotated - -from flytekit.core.artifact import Artifact, Inputs -from flytekit.core.workflow import workflow -from flytekit.trigger import Trigger - - -def test_basic_11(): - # This test translates to - # Trigger(trigger_on=[hourlyArtifact], - # inputs={"x": hourlyArtifact}) - hourlyArtifact = Artifact( - name="hourly_artifact", - time_partitioned=True, - partition_keys=["region"], - ) - aq_idl = hourlyArtifact.embed_as_query([hourlyArtifact]) - assert aq_idl.HasField("binding") - assert aq_idl.binding.index == 0 - assert aq_idl.binding.partition_key == "" - assert aq_idl.binding.bind_to_time_partition is False - - -def test_basic_1(): - # This test would translate to - # Trigger(trigger_on=[hourlyArtifact], - # inputs={"x": hourlyArtifact.query(region="LAX")}) - # note since hourlyArtifact is time partitioned, and it has one other partition key called some_dim, - # these should be bound to the trigger, and region should be a static value. - hourlyArtifact = Artifact( - name="hourly_artifact", - time_partitioned=True, - partition_keys=["region", "some_dim"], - ) - - aq = hourlyArtifact.query(partitions={"region": "LAX"}) - aq_idl = aq.to_flyte_idl([hourlyArtifact]) - assert aq_idl.artifact_id.time_partition.value.HasField("triggered_binding") - assert aq_idl.artifact_id.time_partition.value.triggered_binding.index == 0 - assert aq_idl.artifact_id.time_partition.value.triggered_binding.bind_to_time_partition is True - assert aq_idl.artifact_id.partitions.value["some_dim"].HasField("triggered_binding") - assert aq_idl.artifact_id.partitions.value["some_dim"].triggered_binding.index == 0 - assert aq_idl.artifact_id.partitions.value["some_dim"].triggered_binding.bind_to_time_partition is False - assert aq_idl.artifact_id.partitions.value["region"].static_value == "LAX" - - -def test_basic_2(): - dailyArtifact = Artifact(name="daily_artifact", time_partitioned=True) - - aq = dailyArtifact.query(time_partition=dailyArtifact.time_partition - timedelta(days=1)) - aq_idl = aq.to_flyte_idl([dailyArtifact]) - x = aq_idl.artifact_id.partitions.value - assert len(x) == 0 - assert aq_idl.artifact_id.time_partition.value.triggered_binding.index == 0 - assert aq_idl.artifact_id.time_partition.value.triggered_binding.HasField("partition_key") is False - assert aq_idl.artifact_id.time_partition.value.triggered_binding.bind_to_time_partition is True - assert aq_idl.artifact_id.time_partition.value.triggered_binding.transform is not None - - -def test_big_trigger(): - dailyArtifact = Artifact(name="daily_artifact", time_partitioned=True) - hourlyArtifact = Artifact( - name="hourly_artifact", - time_partitioned=True, - partition_keys=["region"], - ) - UnrelatedArtifact = Artifact(name="unrelated_artifact", time_partitioned=True) - UnrelatedTwo = Artifact(name="unrelated_two", partition_keys=["region"]) - - t = Trigger( - # store these locally. - trigger_on=[dailyArtifact, hourlyArtifact], - inputs={ - # this needs to be serialized into a query. - "today_upstream": dailyArtifact, # this means: use the matched artifact - "yesterday_upstream": dailyArtifact.query(time_partition=dailyArtifact.time_partition - timedelta(days=1)), - # use the matched hourly artifact's time partition, but query on region = LAX - # This is tricky because it's partially bound. - "other_daily_upstream": hourlyArtifact.query(partitions={"region": "LAX"}), - "region": "SEA", # static value that will be passed as input - "other_artifact": UnrelatedArtifact.query(time_partition=dailyArtifact.time_partition), - "other_artifact_2": UnrelatedArtifact.query(time_partition=hourlyArtifact.time_partition), - "other_artifact_3": UnrelatedTwo.query(partitions={"rgg": hourlyArtifact.partitions.region}), - }, - ) - - @workflow - def my_workflow( - today_upstream: str, - yesterday_upstream: str, - other_daily_upstream: str, - region: str, - other_artifact: str, - other_artifact_2: str, - other_artifact_3: str, - dt: datetime, - ) -> Annotated[str, dailyArtifact(time_partition=Inputs.dt)]: - ... - - pm = t.get_parameter_map(my_workflow.python_interface.inputs, my_workflow.interface.inputs) - - assert pm.parameters["today_upstream"].artifact_query == art_id.ArtifactQuery( - binding=art_id.ArtifactBindingData( - index=0, - ), - ) - assert not pm.parameters["today_upstream"].artifact_query.binding.partition_key - assert not pm.parameters["today_upstream"].artifact_query.binding.bind_to_time_partition - assert not pm.parameters["today_upstream"].artifact_query.binding.transform - - assert pm.parameters["yesterday_upstream"].artifact_query == art_id.ArtifactQuery( - artifact_id=art_id.ArtifactID( - artifact_key=art_id.ArtifactKey(project=None, domain=None, name="daily_artifact"), - time_partition=art_id.TimePartition( - value=art_id.LabelValue( - triggered_binding=art_id.ArtifactBindingData(index=0, bind_to_time_partition=True, transform="-P1D") - ), - ), - ), - ) - - assert pm.parameters["other_daily_upstream"].artifact_query == art_id.ArtifactQuery( - artifact_id=art_id.ArtifactID( - artifact_key=art_id.ArtifactKey(project=None, domain=None, name="hourly_artifact"), - time_partition=art_id.TimePartition( - value=art_id.LabelValue( - triggered_binding=art_id.ArtifactBindingData(index=1, bind_to_time_partition=True) - ) - ), - partitions=art_id.Partitions( - value={ - "region": art_id.LabelValue(static_value="LAX"), - } - ), - ), - ) - - assert pm.parameters["region"].default == literals_pb2.Literal( - scalar=literals_pb2.Scalar(primitive=literals_pb2.Primitive(string_value="SEA")) - ) - - assert pm.parameters["other_artifact"].artifact_query == art_id.ArtifactQuery( - artifact_id=art_id.ArtifactID( - artifact_key=art_id.ArtifactKey(project=None, domain=None, name="unrelated_artifact"), - time_partition=art_id.TimePartition( - value=art_id.LabelValue( - triggered_binding=art_id.ArtifactBindingData(index=0, bind_to_time_partition=True) - ) - ), - ) - ) - - assert pm.parameters["other_artifact_2"].artifact_query == art_id.ArtifactQuery( - artifact_id=art_id.ArtifactID( - artifact_key=art_id.ArtifactKey(project=None, domain=None, name="unrelated_artifact"), - time_partition=art_id.TimePartition( - value=art_id.LabelValue( - triggered_binding=art_id.ArtifactBindingData(index=1, bind_to_time_partition=True) - ) - ), - ) - ) - - assert pm.parameters["other_artifact_3"].artifact_query == art_id.ArtifactQuery( - artifact_id=art_id.ArtifactID( - artifact_key=art_id.ArtifactKey(project=None, domain=None, name="unrelated_two"), - partitions=art_id.Partitions( - value={ - "rgg": art_id.LabelValue( - triggered_binding=art_id.ArtifactBindingData(index=1, partition_key="region") - ), - } - ), - ) - ) - - idl_t = t.to_flyte_idl() - assert idl_t.triggers[0].HasField("time_partition") - assert idl_t.triggers[1].HasField("time_partition") - - # Test calling it to create the LaunchPlan object which adds to the global context - @t - @workflow - def tst_wf( - today_upstream: str, - yesterday_upstream: str, - other_daily_upstream: str, - region: str, - other_artifact: str, - other_artifact_2: str, - other_artifact_3: str, - ) -> str: - ... - - -def test_partition_only(): - # This test is different than the tests above in that we're using the value of a partition directly (as opposed to - # running a query parameterized by those partition values). The only case that we'll allow for this is the - # time partition for now. - dailyArtifact = Artifact(name="daily_artifact", time_partitioned=True) - - t = Trigger( - # store these locally. - trigger_on=[dailyArtifact], - inputs={ - "today_upstream": dailyArtifact.time_partition - timedelta(days=1), - }, - ) - - @workflow - def tst_wf( - today_upstream: str, - ): - ... - - pm = t.get_parameter_map(tst_wf.python_interface.inputs, tst_wf.interface.inputs) - assert pm.parameters["today_upstream"].artifact_query == art_id.ArtifactQuery( - binding=art_id.ArtifactBindingData( - index=0, - bind_to_time_partition=True, - transform="-P1D", - ) - ) diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 7eed167246..6673c2fb65 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -22,6 +22,7 @@ from marshmallow_enum import LoadDumpOptions from marshmallow_jsonschema import JSONSchema from mashumaro.mixins.json import DataClassJSONMixin +from mashumaro.mixins.orjson import DataClassORJSONMixin from typing_extensions import Annotated, get_args, get_origin from flytekit import kwtypes @@ -2366,6 +2367,10 @@ def test_DataclassTransformer_get_literal_type(): class MyDataClassMashumaro(DataClassJsonMixin): x: int + @dataclass + class MyDataClassMashumaroORJSON(DataClassJsonMixin): + x: int + @dataclass_json @dataclass class MyDataClass: @@ -2379,6 +2384,9 @@ class MyDataClass: literal_type = de.get_literal_type(MyDataClassMashumaro) assert literal_type is not None + literal_type = de.get_literal_type(MyDataClassMashumaroORJSON) + assert literal_type is not None + invalid_json_str = "{ unbalanced_braces" with pytest.raises(Exception): Literal(scalar=Scalar(generic=_json_format.Parse(invalid_json_str, _struct.Struct()))) @@ -2389,6 +2397,10 @@ def test_DataclassTransformer_to_literal(): class MyDataClassMashumaro(DataClassJsonMixin): x: int + @dataclass + class MyDataClassMashumaroORJSON(DataClassORJSONMixin): + x: int + @dataclass_json @dataclass class MyDataClass: @@ -2398,12 +2410,19 @@ class MyDataClass: ctx = FlyteContext.current_context() my_dat_class_mashumaro = MyDataClassMashumaro(5) + my_dat_class_mashumaro_orjson = MyDataClassMashumaroORJSON(5) my_data_class = MyDataClass(5) lv_mashumaro = transformer.to_literal(ctx, my_dat_class_mashumaro, MyDataClassMashumaro, MyDataClassMashumaro) assert lv_mashumaro is not None assert lv_mashumaro.scalar.generic["x"] == 5 + lv_mashumaro_orjson = transformer.to_literal( + ctx, my_dat_class_mashumaro_orjson, MyDataClassMashumaroORJSON, MyDataClassMashumaroORJSON + ) + assert lv_mashumaro_orjson is not None + assert lv_mashumaro_orjson.scalar.generic["x"] == 5 + lv = transformer.to_literal(ctx, my_data_class, MyDataClass, MyDataClass) assert lv is not None assert lv.scalar.generic["x"] == 5 @@ -2414,6 +2433,10 @@ def test_DataclassTransformer_to_python_value(): class MyDataClassMashumaro(DataClassJsonMixin): x: int + @dataclass + class MyDataClassMashumaroORJSON(DataClassORJSONMixin): + x: int + @dataclass_json @dataclass class MyDataClass: @@ -2432,8 +2455,18 @@ class MyDataClass: assert isinstance(result, MyDataClassMashumaro) assert result.x == 5 + result = de.to_python_value(FlyteContext.current_context(), mock_literal, MyDataClassMashumaroORJSON) + assert isinstance(result, MyDataClassMashumaroORJSON) + assert result.x == 5 + def test_DataclassTransformer_guess_python_type(): + @dataclass + class DatumMashumaroORJSON(DataClassORJSONMixin): + x: int + y: Color + z: datetime.datetime + @dataclass class DatumMashumaro(DataClassJSONMixin): x: int @@ -2464,6 +2497,16 @@ class Datum(DataClassJSONMixin): assert datum_mashumaro.x == pv.x assert datum_mashumaro.y.value == pv.y + lt = TypeEngine.to_literal_type(DatumMashumaroORJSON) + now = datetime.datetime.now() + datum_mashumaro_orjson = DatumMashumaroORJSON(5, Color.RED, now) + lv = transformer.to_literal(ctx, datum_mashumaro_orjson, DatumMashumaroORJSON, lt) + gt = transformer.guess_python_type(lt) + pv = transformer.to_python_value(ctx, lv, expected_python_type=gt) + assert datum_mashumaro_orjson.x == pv.x + assert datum_mashumaro_orjson.y.value == pv.y + assert datum_mashumaro_orjson.z.isoformat() == pv.z + def test_ListTransformer_get_sub_type(): assert ListTransformer.get_sub_type_or_none(typing.List[str]) is str diff --git a/tests/flytekit/unit/core/test_type_hints.py b/tests/flytekit/unit/core/test_type_hints.py index bdbacf246e..a1290d54c0 100644 --- a/tests/flytekit/unit/core/test_type_hints.py +++ b/tests/flytekit/unit/core/test_type_hints.py @@ -32,6 +32,7 @@ from flytekit.core.testing import patch, task_mock from flytekit.core.type_engine import RestrictedTypeError, SimpleTransformer, TypeEngine from flytekit.core.workflow import workflow +from flytekit.exceptions.user import FlyteValidationException from flytekit.models import literals as _literal_models from flytekit.models.core import types as _core_types from flytekit.models.interface import Parameter @@ -131,6 +132,15 @@ def my_task() -> str: assert context_manager.FlyteContextManager.size() == 1 +def test_missing_output(): + @workflow + def wf() -> str: + return None # type: ignore + + with pytest.raises(FlyteValidationException, match="Failed to bind output"): + wf.compile() + + def test_engine_file_output(): basic_blob_type = _core_types.BlobType( format="", @@ -802,7 +812,7 @@ def t1(a: int) -> typing.NamedTuple("OutputsBC", t1_int_output=int, c=str): def t2(a: str) -> str: return a - with pytest.raises(TypeError): + with pytest.raises(FlyteValidationException): @workflow def my_wf(a: int, b: str) -> (int, str): diff --git a/tests/flytekit/unit/core/test_utils.py b/tests/flytekit/unit/core/test_utils.py index 910156f50a..ca0d07565d 100644 --- a/tests/flytekit/unit/core/test_utils.py +++ b/tests/flytekit/unit/core/test_utils.py @@ -1,3 +1,5 @@ +from collections import OrderedDict + import pytest import flytekit @@ -90,7 +92,7 @@ def t() -> str: assert t() == "hello world" assert t.get_config(settings=ss) == {} - ts = get_serializable_task(ss, t) + ts = get_serializable_task(OrderedDict(), ss, t) assert ts.template.config == {"foo": "bar"} @task @@ -98,5 +100,5 @@ def t() -> str: def t() -> str: return "hello world" - ts = get_serializable_task(ss, t) + ts = get_serializable_task(OrderedDict(), ss, t) assert ts.template.config == {"foo": "baz"} diff --git a/tests/flytekit/unit/exceptions/test_scopes.py b/tests/flytekit/unit/exceptions/test_scopes.py index 75ef74383e..0b60d5af95 100644 --- a/tests/flytekit/unit/exceptions/test_scopes.py +++ b/tests/flytekit/unit/exceptions/test_scopes.py @@ -50,6 +50,7 @@ def test_intercepted_scope_non_flyte_exception(): assert e.value == value_error assert "Bad value" in e.verbose_message assert "User error." in e.verbose_message + assert "ValueError:" in e.verbose_message assert e.error_code == "USER:Unknown" assert e.kind == _error_models.ContainerError.Kind.NON_RECOVERABLE @@ -60,6 +61,7 @@ def test_intercepted_scope_non_flyte_exception(): assert e.value == value_error assert "Bad value" in e.verbose_message assert "SYSTEM ERROR!" in e.verbose_message + assert "ValueError:" in e.verbose_message assert e.error_code == "SYSTEM:Unknown" assert e.kind == _error_models.ContainerError.Kind.RECOVERABLE @@ -74,6 +76,7 @@ def test_intercepted_scope_flyte_user_exception(): assert e.value == assertion_error assert "Bad assert" in e.verbose_message assert "User error." in e.verbose_message + assert "FlyteAssertion:" in e.verbose_message assert e.error_code == "USER:AssertionError" assert e.kind == _error_models.ContainerError.Kind.NON_RECOVERABLE @@ -84,6 +87,7 @@ def test_intercepted_scope_flyte_user_exception(): assert e.value == assertion_error assert "Bad assert" in e.verbose_message assert "User error." in e.verbose_message + assert "FlyteAssertion:" in e.verbose_message assert e.error_code == "USER:AssertionError" assert e.kind == _error_models.ContainerError.Kind.NON_RECOVERABLE @@ -98,6 +102,7 @@ def test_intercepted_scope_flyte_system_exception(): assert e.value == assertion_error assert "Bad assert" in e.verbose_message assert "SYSTEM ERROR!" in e.verbose_message + assert "FlyteSystemAssertion:" in e.verbose_message assert e.kind == _error_models.ContainerError.Kind.RECOVERABLE assert e.error_code == "SYSTEM:AssertionError" @@ -108,5 +113,6 @@ def test_intercepted_scope_flyte_system_exception(): assert e.value == assertion_error assert "Bad assert" in e.verbose_message assert "SYSTEM ERROR!" in e.verbose_message + assert "FlyteSystemAssertion:" in e.verbose_message assert e.error_code == "SYSTEM:AssertionError" assert e.kind == _error_models.ContainerError.Kind.RECOVERABLE diff --git a/tests/flytekit/unit/experimental/test_eager_workflows.py b/tests/flytekit/unit/experimental/test_eager_workflows.py index d4438f91f7..9760f8c008 100644 --- a/tests/flytekit/unit/experimental/test_eager_workflows.py +++ b/tests/flytekit/unit/experimental/test_eager_workflows.py @@ -9,7 +9,7 @@ from hypothesis import given, settings from flytekit import dynamic, task, workflow -from flytekit.core.type_engine import TypeTransformerFailedError +from flytekit.exceptions.user import FlyteValidationException from flytekit.experimental import EagerException, eager from flytekit.types.directory import FlyteDirectory from flytekit.types.file import FlyteFile @@ -213,7 +213,7 @@ async def eager_wf(x: int) -> int: out = await local_wf(x=x) return await double(x=out) - with pytest.raises(TypeTransformerFailedError): + with pytest.raises(FlyteValidationException): asyncio.run(eager_wf(x=x_input)) diff --git a/tests/flytekit/unit/extend/test_agent.py b/tests/flytekit/unit/extend/test_agent.py index 0010bccb1f..69f9712dbb 100644 --- a/tests/flytekit/unit/extend/test_agent.py +++ b/tests/flytekit/unit/extend/test_agent.py @@ -18,12 +18,14 @@ DeleteTaskResponse, GetTaskRequest, GetTaskResponse, + ListAgentsRequest, + ListAgentsResponse, Resource, ) from flytekit import PythonFunctionTask, task from flytekit.configuration import FastSerializationSettings, Image, ImageConfig, SerializationSettings -from flytekit.extend.backend.agent_service import AsyncAgentService +from flytekit.extend.backend.agent_service import AgentMetadataService, AsyncAgentService from flytekit.extend.backend.base_agent import ( AgentBase, AgentRegistry, @@ -49,6 +51,8 @@ class Metadata: class DummyAgent(AgentBase): + name = "Dummy Agent" + def __init__(self): super().__init__(task_type="dummy", asynchronous=False) @@ -71,6 +75,8 @@ def delete(self, context: grpc.ServicerContext, resource_meta: bytes) -> DeleteT class AsyncDummyAgent(AgentBase): + name = "Async Dummy Agent" + def __init__(self): super().__init__(task_type="async_dummy", asynchronous=True) @@ -91,6 +97,8 @@ async def async_delete(self, context: grpc.ServicerContext, resource_meta: bytes class SyncDummyAgent(AgentBase): + name = "Sync Dummy Agent" + def __init__(self): super().__init__(task_type="sync_dummy", asynchronous=True) @@ -161,6 +169,10 @@ def __init__(self, **kwargs): with pytest.raises(Exception, match="Cannot find agent for task type: non-exist-type."): t.execute() + agent_metadata = AgentRegistry.get_agent_metadata("Dummy Agent") + assert agent_metadata.name == "Dummy Agent" + assert agent_metadata.supported_task_types == ["dummy"] + @pytest.mark.asyncio async def test_async_dummy_agent(): @@ -175,6 +187,24 @@ async def test_async_dummy_agent(): res = await agent.async_delete(ctx, metadata_bytes) assert res == DeleteTaskResponse() + agent_metadata = AgentRegistry.get_agent_metadata("Async Dummy Agent") + assert agent_metadata.name == "Async Dummy Agent" + assert agent_metadata.supported_task_types == ["async_dummy"] + + +@pytest.mark.asyncio +async def test_sync_dummy_agent(): + AgentRegistry.register(SyncDummyAgent()) + ctx = MagicMock(spec=grpc.ServicerContext) + agent = AgentRegistry.get_agent("sync_dummy") + res = await agent.async_create(ctx, "/tmp", sync_dummy_template, task_inputs) + assert res.resource.state == SUCCEEDED + assert res.resource.outputs == LiteralMap({}).to_flyte_idl() + + agent_metadata = AgentRegistry.get_agent_metadata("Sync Dummy Agent") + assert agent_metadata.name == "Sync Dummy Agent" + assert agent_metadata.supported_task_types == ["sync_dummy"] + @pytest.mark.asyncio async def test_sync_dummy_agent(): @@ -223,6 +253,10 @@ async def run_agent_server(): res = await service.GetTask(GetTaskRequest(task_type=fake_agent, resource_meta=metadata_bytes), ctx) assert res is None + metadata_service = AgentMetadataService() + res = await metadata_service.ListAgent(ListAgentsRequest(), ctx) + assert isinstance(res, ListAgentsResponse) + def test_agent_server(): loop.run_in_executor(None, run_agent_server) diff --git a/tests/flytekit/unit/remote/test_remote.py b/tests/flytekit/unit/remote/test_remote.py index 6a869f4168..e7c296191a 100644 --- a/tests/flytekit/unit/remote/test_remote.py +++ b/tests/flytekit/unit/remote/test_remote.py @@ -207,6 +207,22 @@ def test_more_stuff(mock_client): assert computed_v2 != computed_v3 +@patch("flytekit.remote.remote.SynchronousFlyteClient") +def test_version_hash_special_characters(mock_client): + r = FlyteRemote(config=Config.auto(), default_project="project", default_domain="domain") + + serialization_settings = flytekit.configuration.SerializationSettings( + project="project", + domain="domain", + version="version", + env=None, + image_config=ImageConfig.auto(img_name=DefaultImages.default_image()), + ) + + computed_v = r._version_from_hash(b"", serialization_settings) + assert "=" not in computed_v + + def test_get_extra_headers_azure_blob_storage(): native_url = "abfs://flyte@storageaccount/container/path/to/file" headers = FlyteRemote.get_extra_headers_for_protocol(native_url) diff --git a/tests/flytekit/unit/tools/test_fast_registration.py b/tests/flytekit/unit/tools/test_fast_registration.py index aae3995bcb..8d1c565a9d 100644 --- a/tests/flytekit/unit/tools/test_fast_registration.py +++ b/tests/flytekit/unit/tools/test_fast_registration.py @@ -46,8 +46,7 @@ def flyte_project(tmp_path): def test_package(flyte_project, tmp_path): archive_fname = fast_package(source=flyte_project, output_dir=tmp_path) with tarfile.open(archive_fname) as tar: - assert tar.getnames() == [ - "", # tar root, output removes leading '/' + assert sorted(tar.getnames()) == [ ".dockerignore", ".gitignore", "keep.foo", @@ -67,8 +66,7 @@ def test_package(flyte_project, tmp_path): def test_package_with_symlink(flyte_project, tmp_path): archive_fname = fast_package(source=flyte_project / "src", output_dir=tmp_path, deref_symlinks=True) with tarfile.open(archive_fname, dereference=True) as tar: - assert tar.getnames() == [ - "", # tar root, output removes leading '/' + assert sorted(tar.getnames()) == [ "util", "workflows", "workflows/hello_world.py",