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 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 test_underscore_execute_fall_back_remote_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
options = Options(
raw_output_data_config=common_models.RawOutputDataConfig(
output_location_prefix="raw_output"),
security_context=security.SecurityContext(run_as=security.Identity(
iam_role="iam:some:role")),
)
def local_assertions(*args, **kwargs):
execution_spec = args[3]
assert execution_spec.security_context.run_as.iam_role == "iam:some:role"
assert execution_spec.raw_output_data_config.output_location_prefix == "raw_output"
mock_client.create_execution.side_effect = local_assertions
mock_entity = MagicMock()
remote._execute(
mock_entity,
inputs={},
project="proj",
domain="dev",
options=options,
)
def create_launch_plan(
self,
default_inputs: Dict[str, _promise.Input] = None,
fixed_inputs: Dict[str, Any] = None,
schedule: _schedule_models.Schedule = None,
role: str = None,
notifications: List[_common_models.Notification] = None,
labels: _common_models.Labels = None,
annotations: _common_models.Annotations = None,
assumable_iam_role: str = None,
kubernetes_service_account: str = None,
raw_output_data_prefix: str = None,
):
"""
This method will create a launch plan object that can execute this workflow.
:param dict[Text,flytekit.common.promise.Input] default_inputs:
:param dict[Text,T] fixed_inputs:
:param flytekit.models.schedule.Schedule schedule: A schedule on which to execute this launch plan.
:param Text role: Deprecated. Use assumable_iam_role instead.
:param list[flytekit.models.common.Notification] notifications: A list of notifications to enact by default for
this launch plan.
:param flytekit.models.common.Labels labels:
:param flytekit.models.common.Annotations annotations:
:param cls: This parameter can be used by users to define an extension of a launch plan to instantiate. The
class provided should be a subclass of flytekit.common.launch_plan.SdkLaunchPlan.
:param Text assumable_iam_role: The IAM role to execute the workflow with.
:param Text kubernetes_service_account: The kubernetes service account to execute the workflow with.
:param Text raw_output_data_prefix: Bucket for offloaded data
:rtype: flytekit.common.launch_plan.SdkRunnableLaunchPlan
"""
# TODO: Actually ensure the parameters conform.
if role and (assumable_iam_role or kubernetes_service_account):
raise ValueError("Cannot set both role and auth. Role is deprecated, use auth instead.")
fixed_inputs = fixed_inputs or {}
merged_default_inputs = {v.name: v for v in self.user_inputs if v.name not in fixed_inputs}
merged_default_inputs.update(default_inputs or {})
if role:
assumable_iam_role = role # For backwards compatibility
auth_role = _common_models.AuthRole(
assumable_iam_role=assumable_iam_role, kubernetes_service_account=kubernetes_service_account,
)
raw_output_config = _common_models.RawOutputDataConfig(raw_output_data_prefix or "")
return _launch_plan.SdkRunnableLaunchPlan(
sdk_workflow=self,
default_inputs={
k: user_input.rename_and_return_reference(k) for k, user_input in _six.iteritems(merged_default_inputs)
},
fixed_inputs=fixed_inputs,
schedule=schedule,
notifications=notifications,
labels=labels,
annotations=annotations,
auth_role=auth_role,
raw_output_data_config=raw_output_config,
)
def raw_output_data_config(self):
"""
:rtype: flytekit.models.common.RawOutputDataConfig
"""
raw_output_data_config = super(SdkLaunchPlan, self).raw_output_data_config
if raw_output_data_config is not None and raw_output_data_config.output_location_prefix != "":
return raw_output_data_config
# If it was not set explicitly then let's use the value found in the configuration.
return _common_models.RawOutputDataConfig(_auth_config.RAW_OUTPUT_DATA_PREFIX.get())
def default_from(
cls,
k8s_service_account: typing.Optional[str] = None,
raw_data_prefix: typing.Optional[str] = None) -> "Options":
return cls(
security_context=security.SecurityContext(run_as=security.Identity(
k8s_service_account=k8s_service_account))
if k8s_service_account else None,
raw_output_data_config=common_models.RawOutputDataConfig(
output_location_prefix=raw_data_prefix)
if raw_data_prefix else None,
)
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_raw_output_data_config():
obj = _common.RawOutputDataConfig("s3://bucket")
assert obj.output_location_prefix == "s3://bucket"
obj2 = _common.RawOutputDataConfig.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 test_execution_spec(literal_value_pair):
literal_value, _ = literal_value_pair
obj = _execution.ExecutionSpec(
_identifier.Identifier(_identifier.ResourceType.LAUNCH_PLAN, "project", "domain", "name", "version"),
_execution.ExecutionMetadata(_execution.ExecutionMetadata.ExecutionMode.MANUAL, "tester", 1),
notifications=_execution.NotificationList(
[
_common_models.Notification(
[_core_exec.WorkflowExecutionPhase.ABORTED],
pager_duty=_common_models.PagerDutyNotification(recipients_email=["a", "b", "c"]),
)
]
),
raw_output_data_config=_common_models.RawOutputDataConfig(output_location_prefix="raw_output"),
max_parallelism=100,
)
assert obj.launch_plan.resource_type == _identifier.ResourceType.LAUNCH_PLAN
assert obj.launch_plan.domain == "domain"
assert obj.launch_plan.project == "project"
assert obj.launch_plan.name == "name"
assert obj.launch_plan.version == "version"
assert obj.metadata.mode == _execution.ExecutionMetadata.ExecutionMode.MANUAL
assert obj.metadata.nesting == 1
assert obj.metadata.principal == "tester"
assert obj.notifications.notifications[0].phases == [_core_exec.WorkflowExecutionPhase.ABORTED]
assert obj.notifications.notifications[0].pager_duty.recipients_email == [
"a",
"b",
"c",
]
assert obj.disable_all is None
assert obj.max_parallelism == 100
assert obj.raw_output_data_config.output_location_prefix == "raw_output"
obj2 = _execution.ExecutionSpec.from_flyte_idl(obj.to_flyte_idl())
assert obj == obj2
assert obj2.launch_plan.resource_type == _identifier.ResourceType.LAUNCH_PLAN
assert obj2.launch_plan.domain == "domain"
assert obj2.launch_plan.project == "project"
assert obj2.launch_plan.name == "name"
assert obj2.launch_plan.version == "version"
assert obj2.metadata.mode == _execution.ExecutionMetadata.ExecutionMode.MANUAL
assert obj2.metadata.nesting == 1
assert obj2.metadata.principal == "tester"
assert obj2.notifications.notifications[0].phases == [_core_exec.WorkflowExecutionPhase.ABORTED]
assert obj2.notifications.notifications[0].pager_duty.recipients_email == [
"a",
"b",
"c",
]
assert obj2.disable_all is None
assert obj2.max_parallelism == 100
assert obj2.raw_output_data_config.output_location_prefix == "raw_output"
obj = _execution.ExecutionSpec(
_identifier.Identifier(_identifier.ResourceType.LAUNCH_PLAN, "project", "domain", "name", "version"),
_execution.ExecutionMetadata(_execution.ExecutionMetadata.ExecutionMode.MANUAL, "tester", 1),
disable_all=True,
)
assert obj.launch_plan.resource_type == _identifier.ResourceType.LAUNCH_PLAN
assert obj.launch_plan.domain == "domain"
assert obj.launch_plan.project == "project"
assert obj.launch_plan.name == "name"
assert obj.launch_plan.version == "version"
assert obj.metadata.mode == _execution.ExecutionMetadata.ExecutionMode.MANUAL
assert obj.metadata.nesting == 1
assert obj.metadata.principal == "tester"
assert obj.notifications is None
assert obj.disable_all is True
obj2 = _execution.ExecutionSpec.from_flyte_idl(obj.to_flyte_idl())
assert obj == obj2
assert obj2.launch_plan.resource_type == _identifier.ResourceType.LAUNCH_PLAN
assert obj2.launch_plan.domain == "domain"
assert obj2.launch_plan.project == "project"
assert obj2.launch_plan.name == "name"
assert obj2.launch_plan.version == "version"
assert obj2.metadata.mode == _execution.ExecutionMetadata.ExecutionMode.MANUAL
assert obj2.metadata.nesting == 1
assert obj2.metadata.principal == "tester"
assert obj2.notifications is None
assert obj2.disable_all is True
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
