diff --git a/flytekit/core/map_task.py b/flytekit/core/map_task.py index 29838cffcd..8af351f6a8 100644 --- a/flytekit/core/map_task.py +++ b/flytekit/core/map_task.py @@ -5,20 +5,42 @@ import os from contextlib import contextmanager +from dataclasses import dataclass from itertools import count from typing import Any, Dict, List, Optional, Type +from dataclasses_json import dataclass_json + from flytekit.core.base_task import PythonTask from flytekit.core.constants import SdkTaskType from flytekit.core.context_manager import ExecutionState, FlyteContext, FlyteContextManager, SerializationSettings from flytekit.core.interface import transform_interface_to_list_interface from flytekit.core.python_function_task import PythonFunctionTask from flytekit.exceptions import scopes as exception_scopes -from flytekit.models.array_job import ArrayJob from flytekit.models.interface import Variable from flytekit.models.task import Container, K8sPod, Sql +@dataclass_json +@dataclass +class ArrayJob: + """ + Initializes a new ArrayJob. + :param int parallelism: Defines the minimum number of instances to bring up concurrently at any given point. + :param int size: Defines the number of instances to launch at most. This number should match the size of + the input if the job requires processing of all input data. This has to be a positive number. + :param int min_successes: An absolute number of the minimum number of successful completions of subtasks. As + soon as this criteria is met, the array job will be marked as successful and outputs will be computed. + :param float min_success_ratio: Determines the minimum fraction of total jobs which can complete successfully + before terminating the job and marking it successful. + """ + + parallelism: Optional[int] = None + size: Optional[int] = None + min_successes: Optional[int] = None + min_success_ratio: Optional[float] = None + + class MapPythonTask(PythonTask): """ A MapPythonTask defines a :py:class:`flytekit.PythonTask` which specifies how to run @@ -109,11 +131,11 @@ def get_sql(self, settings: SerializationSettings) -> Sql: with self.prepare_target(): return self._run_task.get_sql(settings) - def get_custom(self, settings: SerializationSettings) -> Dict[str, Any]: - return ArrayJob(parallelism=self._max_concurrency, min_success_ratio=self._min_success_ratio).to_dict() - def get_config(self, settings: SerializationSettings) -> Dict[str, str]: - return self._run_task.get_config(settings) + array_job = ArrayJob(parallelism=self._max_concurrency, min_success_ratio=self._min_success_ratio).to_dict() + if self._run_task.get_config(settings) is not None: + array_job.update(self._run_task.get_config(settings)) + return {str(key): str(value) for key, value in array_job.items()} @property def run_task(self) -> PythonTask: diff --git a/flytekit/models/array_job.py b/flytekit/models/array_job.py index 4e4bf99cc7..46e5d60371 100644 --- a/flytekit/models/array_job.py +++ b/flytekit/models/array_job.py @@ -1,11 +1,16 @@ import json as _json +from deprecated import deprecated as _deprecated from flyteidl.plugins import array_job_pb2 as _array_job from google.protobuf import json_format as _json_format from flytekit.models import common as _common +@_deprecated( + reason="Use ArrayJob in flytekit.core.map_task instead", + version="0.25.0", +) class ArrayJob(_common.FlyteCustomIdlEntity): def __init__(self, parallelism=None, size=None, min_successes=None, min_success_ratio=None): """ diff --git a/plugins/flytekit-k8s-pod/tests/test_pod.py b/plugins/flytekit-k8s-pod/tests/test_pod.py index 82a0bcf6f8..ca4bcb2cfc 100644 --- a/plugins/flytekit-k8s-pod/tests/test_pod.py +++ b/plugins/flytekit-k8s-pod/tests/test_pod.py @@ -329,7 +329,13 @@ def simple_pod_task(i: int): "task-name", "simple_pod_task", ] - assert {"primary_container_name": "primary"} == mapped_task.get_config(serialization_settings) + assert { + "min_success_ratio": "None", + "min_successes": "None", + "parallelism": "None", + "primary_container_name": "primary", + "size": "None", + } == mapped_task.get_config(serialization_settings) def test_fast_pod_task_serialization(): diff --git a/tests/flytekit/unit/models/test_dynamic_job.py b/tests/flytekit/unit/models/test_dynamic_job.py index 1aff800abd..9ea5f4538f 100644 --- a/tests/flytekit/unit/models/test_dynamic_job.py +++ b/tests/flytekit/unit/models/test_dynamic_job.py @@ -4,7 +4,7 @@ import pytest from google.protobuf import text_format -from flytekit.models import array_job as _array_job +from flytekit.core.map_task import ArrayJob from flytekit.models import dynamic_job as _dynamic_job from flytekit.models import literals as _literals from flytekit.models import task as _task @@ -18,7 +18,7 @@ "python", task_metadata, interfaces, - _array_job.ArrayJob(2, 2, 2).to_dict(), + ArrayJob(2, 2, 2).to_dict(), container=_task.Container( "my_image", ["this", "is", "a", "cmd"],