def get_serializable_launch_plan(settings: SerializationSettings,
entity: FlyteLocalEntity,
fast: bool) -> FlyteControlPlaneEntity:
sdk_workflow = get_serializable(settings, entity.workflow)
cp_entity = SdkLaunchPlan(
workflow_id=sdk_workflow.id,
entity_metadata=_launch_plan_models.LaunchPlanMetadata(
schedule=entity.schedule,
notifications=entity.notifications,
),
default_inputs=entity.parameters,
fixed_inputs=entity.fixed_inputs,
labels=entity.labels or _common_models.Labels({}),
annotations=entity.annotations or _common_models.Annotations({}),
auth_role=entity._auth_role,
raw_output_data_config=entity.raw_output_data_config,
)
# These two things are normally set to None in the SdkLaunchPlan constructor and filled in by
# SdkRunnableLaunchPlan/the registration process, so we need to set them manually. The reason is because these
# fields are not part of the underlying LaunchPlanSpec
cp_entity._interface = sdk_workflow.interface
cp_entity._id = _identifier_model.Identifier(
resource_type=_identifier_model.ResourceType.LAUNCH_PLAN,
project=settings.project,
domain=settings.domain,
name=entity.name,
version=settings.version,
)
return cp_entity
def __init__(self,
launch_plan,
inputs,
metadata,
notifications=None,
disable_all=None,
labels=None,
annotations=None):
"""
:param flytekit.models.core.identifier.Identifier launch_plan: Launch plan unique identifier to execute
:param flytekit.models.literals.LiteralMap inputs: Inputs to apply to the launch plan
:param ExecutionMetadata metadata: The metadata to be associated with this execution
:param NotificationList notifications: List of notifications for this execution.
:param bool disable_all: If true, all notifications should be disabled.
:param flytekit.models.common.Labels labels: Labels to apply to the execution.
:param flytekit.models.common.Annotations annotations: Annotations to apply to the execution
"""
self._launch_plan = launch_plan
self._inputs = inputs
self._metadata = metadata
self._notifications = notifications
self._disable_all = disable_all
self._labels = labels or _common_models.Labels({})
self._annotations = annotations or _common_models.Annotations({})
def test_underscore_execute_uses_launch_plan_attributes(mock_wf_exec):
mock_wf_exec.return_value = True
mock_client = MagicMock()
remote = FlyteRemote(config=Config.auto(),
default_project="p1",
default_domain="d1")
remote._client = mock_client
def local_assertions(*args, **kwargs):
execution_spec = args[3]
assert execution_spec.security_context.run_as.k8s_service_account == "svc"
assert execution_spec.labels == common_models.Labels(
{"a": "my_label_value"})
assert execution_spec.annotations == common_models.Annotations(
{"b": "my_annotation_value"})
mock_client.create_execution.side_effect = local_assertions
mock_entity = MagicMock()
options = Options(
labels=common_models.Labels({"a": "my_label_value"}),
annotations=common_models.Annotations({"b": "my_annotation_value"}),
security_context=security.SecurityContext(run_as=security.Identity(
k8s_service_account="svc")),
)
remote._execute(
mock_entity,
inputs={},
project="proj",
domain="dev",
options=options,
)
def __init__(
self,
launch_plan,
metadata,
notifications=None,
disable_all=None,
labels=None,
annotations=None,
auth_role=None,
):
"""
:param flytekit.models.core.identifier.Identifier launch_plan: Launch plan unique identifier to execute
:param ExecutionMetadata metadata: The metadata to be associated with this execution
:param NotificationList notifications: List of notifications for this execution.
:param bool disable_all: If true, all notifications should be disabled.
:param flytekit.models.common.Labels labels: Labels to apply to the execution.
:param flytekit.models.common.Annotations annotations: Annotations to apply to the execution
:param flytekit.models.common.AuthRole auth_role: The authorization method with which to execute the workflow.
"""
self._launch_plan = launch_plan
self._metadata = metadata
self._notifications = notifications
self._disable_all = disable_all
self._labels = labels or _common_models.Labels({})
self._annotations = annotations or _common_models.Annotations({})
self._auth_role = auth_role or _common_models.AuthRole()
def create_launch_plan(self, *args, **kwargs):
# TODO: Correct after implementing new launch plan
assumable_iam_role = _auth_config.ASSUMABLE_IAM_ROLE.get()
kubernetes_service_account = _auth_config.KUBERNETES_SERVICE_ACCOUNT.get(
)
if not (assumable_iam_role or kubernetes_service_account):
raise _user_exceptions.FlyteValidationException(
"No assumable role or service account found")
auth_role = _common_models.AuthRole(
assumable_iam_role=assumable_iam_role,
kubernetes_service_account=kubernetes_service_account,
)
return SdkLaunchPlan(
workflow_id=self.id,
entity_metadata=_launch_plan_models.LaunchPlanMetadata(
schedule=_schedule_models.Schedule(""),
notifications=[],
),
default_inputs=_interface_models.ParameterMap({}),
fixed_inputs=_literal_models.LiteralMap(literals={}),
labels=_common_models.Labels({}),
annotations=_common_models.Annotations({}),
auth_role=auth_role,
raw_output_data_config=_common_models.RawOutputDataConfig(""),
)
def local_assertions(*args, **kwargs):
execution_spec = args[3]
assert execution_spec.security_context.run_as.k8s_service_account == "svc"
assert execution_spec.labels == common_models.Labels(
{"a": "my_label_value"})
assert execution_spec.annotations == common_models.Annotations(
{"b": "my_annotation_value"})
def __init__(
self,
launch_plan,
metadata,
notifications=None,
disable_all=None,
labels=None,
annotations=None,
auth_role=None,
max_parallelism=None,
):
"""
:param flytekit.models.core.identifier.Identifier launch_plan: Launch plan unique identifier to execute
:param ExecutionMetadata metadata: The metadata to be associated with this execution
:param NotificationList notifications: List of notifications for this execution.
:param bool disable_all: If true, all notifications should be disabled.
:param flytekit.models.common.Labels labels: Labels to apply to the execution.
:param flytekit.models.common.Annotations annotations: Annotations to apply to the execution
:param flytekit.models.common.AuthRole auth_role: The authorization method with which to execute the workflow.
:param max_parallelism int: Controls the maximum number of tasknodes that can be run in parallel for the entire
workflow. This is useful to achieve fairness. Note: MapTasks are regarded as one unit, and
parallelism/concurrency of MapTasks is independent from this.
"""
self._launch_plan = launch_plan
self._metadata = metadata
self._notifications = notifications
self._disable_all = disable_all
self._labels = labels or _common_models.Labels({})
self._annotations = annotations or _common_models.Annotations({})
self._auth_role = auth_role or _common_models.AuthRole()
self._max_parallelism = max_parallelism
def test_execution_label_overrides(mock_client_factory):
mock_client = MagicMock()
mock_client.create_execution = MagicMock(
return_value=identifier.WorkflowExecutionIdentifier('xp', 'xd', 'xn'))
mock_client_factory.return_value = mock_client
m = MagicMock()
type(m).id = PropertyMock(return_value=identifier.Identifier(
identifier.ResourceType.LAUNCH_PLAN, "project", "domain", "name",
"version"))
labels = _common_models.Labels({"my": "label"})
engine.FlyteLaunchPlan(m).execute('xp',
'xd',
'xn',
literals.LiteralMap({}),
notification_overrides=[],
label_overrides=labels)
mock_client.create_execution.assert_called_once_with(
'xp', 'xd', 'xn',
_execution_models.ExecutionSpec(
identifier.Identifier(identifier.ResourceType.LAUNCH_PLAN,
"project", "domain", "name", "version"),
literals.LiteralMap({}),
_execution_models.ExecutionMetadata(
_execution_models.ExecutionMetadata.ExecutionMode.MANUAL,
'sdk', 0),
disable_all=True,
labels=labels,
))
def test_launch_plan_spec():
identifier_model = identifier.Identifier(identifier.ResourceType.TASK,
"project", "domain", "name",
"version")
s = schedule.Schedule("asdf", "1 3 4 5 6 7")
launch_plan_metadata_model = launch_plan.LaunchPlanMetadata(
schedule=s, notifications=[])
v = interface.Variable(types.LiteralType(simple=types.SimpleType.BOOLEAN),
"asdf asdf asdf")
p = interface.Parameter(var=v)
parameter_map = interface.ParameterMap({"ppp": p})
fixed_inputs = literals.LiteralMap({
"a":
literals.Literal(scalar=literals.Scalar(primitive=literals.Primitive(
integer=1)))
})
labels_model = common.Labels({})
annotations_model = common.Annotations({"my": "annotation"})
auth_role_model = common.AuthRole(assumable_iam_role="my:iam:role")
raw_data_output_config = common.RawOutputDataConfig("s3://bucket")
empty_raw_data_output_config = common.RawOutputDataConfig("")
max_parallelism = 100
lp_spec_raw_output_prefixed = launch_plan.LaunchPlanSpec(
identifier_model,
launch_plan_metadata_model,
parameter_map,
fixed_inputs,
labels_model,
annotations_model,
auth_role_model,
raw_data_output_config,
max_parallelism,
)
obj2 = launch_plan.LaunchPlanSpec.from_flyte_idl(
lp_spec_raw_output_prefixed.to_flyte_idl())
assert obj2 == lp_spec_raw_output_prefixed
lp_spec_no_prefix = launch_plan.LaunchPlanSpec(
identifier_model,
launch_plan_metadata_model,
parameter_map,
fixed_inputs,
labels_model,
annotations_model,
auth_role_model,
empty_raw_data_output_config,
max_parallelism,
)
obj2 = launch_plan.LaunchPlanSpec.from_flyte_idl(
lp_spec_no_prefix.to_flyte_idl())
assert obj2 == lp_spec_no_prefix
def get_serializable_references(
entity_mapping: OrderedDict,
settings: SerializationSettings,
entity: FlyteLocalEntity,
fast: bool,
) -> FlyteControlPlaneEntity:
# TODO: This entire function isn't necessary. We should just return None or raise an Exception or something.
# Reference entities should already exist on the Admin control plane - they should not be serialized/registered
# again. Indeed we don't actually have enough information to serialize it properly.
if isinstance(entity, ReferenceTask):
cp_entity = SdkTask(
type="ignore",
metadata=TaskMetadata().to_taskmetadata_model(),
interface=entity.interface,
custom={},
container=None,
)
elif isinstance(entity, ReferenceWorkflow):
workflow_metadata = WorkflowMetadata(
on_failure=WorkflowFailurePolicy.FAIL_IMMEDIATELY)
cp_entity = SdkWorkflow(
nodes=[], # Fake an empty list for nodes,
id=entity.reference.id,
metadata=workflow_metadata,
metadata_defaults=workflow_model.WorkflowMetadataDefaults(),
interface=entity.interface,
output_bindings=[],
)
elif isinstance(entity, ReferenceLaunchPlan):
cp_entity = SdkLaunchPlan(
workflow_id=None,
entity_metadata=_launch_plan_models.LaunchPlanMetadata(
schedule=None, notifications=[]),
default_inputs=interface_models.ParameterMap({}),
fixed_inputs=literal_models.LiteralMap({}),
labels=_common_models.Labels({}),
annotations=_common_models.Annotations({}),
auth_role=_common_models.AuthRole(assumable_iam_role="fake:role"),
raw_output_data_config=RawOutputDataConfig(""),
)
# Because of how SdkNodes work, it needs one of these interfaces
# Hopefully this is more trickery that can be cleaned up in the future
cp_entity._interface = TypedInterface.promote_from_model(
entity.interface)
else:
raise Exception("Invalid reference type when serializing")
# Make sure we don't serialize this
cp_entity._has_registered = True
cp_entity.assign_name(entity.id.name)
cp_entity._id = entity.id
return cp_entity
def test_labels():
workflow_to_test = _workflow.workflow(
{},
inputs={
'required_input': _workflow.Input(_types.Types.Integer),
'default_input': _workflow.Input(_types.Types.Integer, default=5)
})
lp = workflow_to_test.create_launch_plan(
fixed_inputs={'required_input': 5},
schedule=_schedules.CronSchedule("* * ? * * *"),
role='what',
labels=_common_models.Labels({"my": "label"}))
assert lp.labels.values == {"my": "label"}
def test_old_style_role():
identifier_model = identifier.Identifier(identifier.ResourceType.TASK,
"project", "domain", "name",
"version")
s = schedule.Schedule("asdf", "1 3 4 5 6 7")
launch_plan_metadata_model = launch_plan.LaunchPlanMetadata(
schedule=s, notifications=[])
v = interface.Variable(types.LiteralType(simple=types.SimpleType.BOOLEAN),
"asdf asdf asdf")
p = interface.Parameter(var=v)
parameter_map = interface.ParameterMap({"ppp": p})
fixed_inputs = literals.LiteralMap({
"a":
literals.Literal(scalar=literals.Scalar(primitive=literals.Primitive(
integer=1)))
})
labels_model = common.Labels({})
annotations_model = common.Annotations({"my": "annotation"})
raw_data_output_config = common.RawOutputDataConfig("s3://bucket")
old_role = _launch_plan_idl.Auth(
kubernetes_service_account="my:service:account")
old_style_spec = _launch_plan_idl.LaunchPlanSpec(
workflow_id=identifier_model.to_flyte_idl(),
entity_metadata=launch_plan_metadata_model.to_flyte_idl(),
default_inputs=parameter_map.to_flyte_idl(),
fixed_inputs=fixed_inputs.to_flyte_idl(),
labels=labels_model.to_flyte_idl(),
annotations=annotations_model.to_flyte_idl(),
raw_output_data_config=raw_data_output_config.to_flyte_idl(),
auth=old_role,
)
lp_spec = launch_plan.LaunchPlanSpec.from_flyte_idl(old_style_spec)
assert lp_spec.auth_role.assumable_iam_role == "my:service:account"
def get_serializable_launch_plan(
entity_mapping: OrderedDict,
settings: SerializationSettings,
entity: LaunchPlan,
fast: bool,
) -> _launch_plan_models.LaunchPlan:
wf_spec = get_serializable(entity_mapping, settings, entity.workflow)
lps = _launch_plan_models.LaunchPlanSpec(
workflow_id=wf_spec.template.id,
entity_metadata=_launch_plan_models.LaunchPlanMetadata(
schedule=entity.schedule,
notifications=entity.notifications,
),
default_inputs=entity.parameters,
fixed_inputs=entity.fixed_inputs,
labels=entity.labels or _common_models.Labels({}),
annotations=entity.annotations or _common_models.Annotations({}),
auth_role=entity._auth_role or _common_models.AuthRole(),
raw_output_data_config=entity.raw_output_data_config
or _common_models.RawOutputDataConfig(""),
)
lp_id = _identifier_model.Identifier(
resource_type=_identifier_model.ResourceType.LAUNCH_PLAN,
project=settings.project,
domain=settings.domain,
name=entity.name,
version=settings.version,
)
lp_model = _launch_plan_models.LaunchPlan(
id=lp_id,
spec=lps,
closure=_launch_plan_models.LaunchPlanClosure(
state=None,
expected_inputs=interface_models.ParameterMap({}),
expected_outputs=interface_models.VariableMap({}),
),
)
return lp_model
def test_promote_from_model():
workflow_to_test = _workflow.workflow(
{},
inputs={
'required_input': _workflow.Input(_types.Types.Integer),
'default_input': _workflow.Input(_types.Types.Integer, default=5)
})
workflow_to_test._id = _identifier.Identifier(
_identifier.ResourceType.WORKFLOW, "p", "d", "n", "v")
lp = workflow_to_test.create_launch_plan(
fixed_inputs={'required_input': 5},
schedule=_schedules.CronSchedule("* * ? * * *"),
role='what',
labels=_common_models.Labels({"my": "label"}))
with _pytest.raises(_user_exceptions.FlyteAssertion):
_launch_plan.SdkRunnableLaunchPlan.from_flyte_idl(lp.to_flyte_idl())
lp_from_spec = _launch_plan.SdkLaunchPlan.from_flyte_idl(lp.to_flyte_idl())
assert not isinstance(lp_from_spec, _launch_plan.SdkRunnableLaunchPlan)
assert isinstance(lp_from_spec, _launch_plan.SdkLaunchPlan)
assert lp_from_spec == lp
def test_labels():
obj = _common.Labels({"my": "label"})
assert obj.values == {"my": "label"}
obj2 = _common.Labels.from_flyte_idl(obj.to_flyte_idl())
assert obj2 == obj
def __init__(
self,
sdk_workflow,
default_inputs=None,
fixed_inputs=None,
role=None,
schedule=None,
notifications=None,
labels=None,
annotations=None,
auth_role=None,
raw_output_data_config=None,
):
"""
:param flytekit.common.local_workflow.SdkRunnableWorkflow sdk_workflow:
:param dict[Text,flytekit.common.promise.Input] default_inputs:
:param dict[Text,Any] fixed_inputs: These inputs will be fixed and not need to be set when executing this
launch plan.
:param Text role: Deprecated. IAM role to execute this launch plan with.
:param flytekit.models.schedule.Schedule: Schedule to apply to this workflow.
:param list[flytekit.models.common.Notification]: List of notifications to apply to this launch plan.
:param flytekit.models.common.Labels labels: Any custom kubernetes labels to apply to workflows executed by this
launch plan.
:param flytekit.models.common.Annotations annotations: Any custom kubernetes annotations to apply to workflows
executed by this launch plan.
Any custom kubernetes annotations to apply to workflows executed by this launch plan.
:param flytekit.models.common.Authrole auth_role: The auth method with which to execute the workflow.
:param flytekit.models.common.RawOutputDataConfig raw_output_data_config: Config for offloading data
"""
if role and auth_role:
raise ValueError(
"Cannot set both role and auth. Role is deprecated, use auth instead."
)
fixed_inputs = fixed_inputs or {}
default_inputs = default_inputs or {}
if role:
auth_role = _common_models.AuthRole(assumable_iam_role=role)
# The constructor for SdkLaunchPlan sets the id to None anyways so we don't bother passing in an ID. The ID
# should be set in one of three places,
# 1) When the object is registered (in the code above)
# 2) By the dynamic task code after this runnable object has already been __call__'ed. The SdkNode produced
# maintains a link to this object and will set the ID according to the configuration variables present.
# 3) When SdkLaunchPlan.fetch() is run
super(SdkRunnableLaunchPlan, self).__init__(
None,
_launch_plan_models.LaunchPlanMetadata(
schedule=schedule or _schedule_model.Schedule(""),
notifications=notifications or [],
),
_interface_models.ParameterMap(default_inputs),
_type_helpers.pack_python_std_map_to_literal_map(
fixed_inputs,
{
k: _type_helpers.get_sdk_type_from_literal_type(var.type)
for k, var in _six.iteritems(sdk_workflow.interface.inputs)
if k in fixed_inputs
},
),
labels or _common_models.Labels({}),
annotations or _common_models.Annotations({}),
auth_role,
raw_output_data_config or _common_models.RawOutputDataConfig(""),
)
self._interface = _interface.TypedInterface(
{k: v.var
for k, v in _six.iteritems(default_inputs)},
sdk_workflow.interface.outputs,
)
self._upstream_entities = {sdk_workflow}
self._sdk_workflow = sdk_workflow
def __init__(
self,
sdk_workflow,
default_inputs=None,
fixed_inputs=None,
role=None,
schedule=None,
notifications=None,
labels=None,
annotations=None,
auth=None,
):
"""
:param flytekit.common.workflow.SdkWorkflow sdk_workflow:
:param dict[Text,flytekit.common.promise.Input] default_inputs:
:param dict[Text,Any] fixed_inputs: These inputs will be fixed and not need to be set when executing this
launch plan.
:param Text role: Deprecated. IAM role to execute this launch plan with.
:param flytekit.models.schedule.Schedule: Schedule to apply to this workflow.
:param list[flytekit.models.common.Notification]: List of notifications to apply to this launch plan.
:param flytekit.models.common.Labels labels: Any custom kubernetes labels to apply to workflows executed by this
launch plan.
:param flytekit.models.common.Annotations annotations: Any custom kubernetes annotations to apply to workflows
executed by this launch plan.
Any custom kubernetes annotations to apply to workflows executed by this launch plan.
:param flytekit.models.launch_plan.Auth auth: The auth method with which to execute the workflow.
"""
if role and auth:
raise ValueError(
"Cannot set both role and auth. Role is deprecated, use auth instead."
)
fixed_inputs = fixed_inputs or {}
default_inputs = default_inputs or {}
if role:
auth = _launch_plan_models.Auth(assumable_iam_role=role)
super(SdkRunnableLaunchPlan, self).__init__(
_identifier.Identifier(_identifier_model.ResourceType.WORKFLOW,
_internal_config.PROJECT.get(),
_internal_config.DOMAIN.get(),
_uuid.uuid4().hex,
_internal_config.VERSION.get()),
_launch_plan_models.LaunchPlanMetadata(
schedule=schedule or _schedule_model.Schedule(''),
notifications=notifications or []),
_interface_models.ParameterMap(default_inputs),
_type_helpers.pack_python_std_map_to_literal_map(
fixed_inputs, {
k: _type_helpers.get_sdk_type_from_literal_type(var.type)
for k, var in _six.iteritems(sdk_workflow.interface.inputs)
if k in fixed_inputs
}),
labels or _common_models.Labels({}),
annotations or _common_models.Annotations({}),
auth,
)
self._interface = _interface.TypedInterface(
{k: v.var
for k, v in _six.iteritems(default_inputs)},
sdk_workflow.interface.outputs)
self._upstream_entities = {sdk_workflow}
self._sdk_workflow = sdk_workflow
def get_serializable_launch_plan(
entity_mapping: OrderedDict,
settings: SerializationSettings,
entity: LaunchPlan,
recurse_downstream: bool = True,
options: Optional[Options] = None,
) -> _launch_plan_models.LaunchPlan:
"""
:param entity_mapping:
:param settings:
:param entity:
:param options:
:param recurse_downstream: This boolean indicate is wf for the entity should also be recursed to
:return:
"""
if recurse_downstream:
wf_spec = get_serializable(entity_mapping, settings, entity.workflow,
options)
wf_id = wf_spec.template.id
else:
wf_id = _identifier_model.Identifier(
resource_type=_identifier_model.ResourceType.WORKFLOW,
project=settings.project,
domain=settings.domain,
name=entity.workflow.name,
version=settings.version,
)
if not options:
options = Options()
if options and options.raw_output_data_config:
raw_prefix_config = options.raw_output_data_config
else:
raw_prefix_config = entity.raw_output_data_config or _common_models.RawOutputDataConfig(
"")
lps = _launch_plan_models.LaunchPlanSpec(
workflow_id=wf_id,
entity_metadata=_launch_plan_models.LaunchPlanMetadata(
schedule=entity.schedule,
notifications=options.notifications or entity.notifications,
),
default_inputs=entity.parameters,
fixed_inputs=entity.fixed_inputs,
labels=options.labels or entity.labels or _common_models.Labels({}),
annotations=options.annotations or entity.annotations
or _common_models.Annotations({}),
auth_role=None,
raw_output_data_config=raw_prefix_config,
max_parallelism=options.max_parallelism or entity.max_parallelism,
security_context=options.security_context or entity.security_context,
)
lp_id = _identifier_model.Identifier(
resource_type=_identifier_model.ResourceType.LAUNCH_PLAN,
project=settings.project,
domain=settings.domain,
name=entity.name,
version=settings.version,
)
lp_model = _launch_plan_models.LaunchPlan(
id=lp_id,
spec=lps,
closure=_launch_plan_models.LaunchPlanClosure(
state=None,
expected_inputs=interface_models.ParameterMap({}),
expected_outputs=interface_models.VariableMap({}),
),
)
return lp_model
