def from_flyte_idl(cls, pb2):
"""
:param flyteidl.admin.launch_plan_pb2.LaunchPlanSpec pb2:
:rtype: LaunchPlanSpec
"""
auth_role = None
# First check the newer field, auth_role.
if pb2.auth_role is not None and (
pb2.auth_role.assumable_iam_role
or pb2.auth_role.kubernetes_service_account):
auth_role = _common.AuthRole.from_flyte_idl(pb2.auth_role)
# Fallback to the deprecated field.
elif pb2.auth is not None:
if pb2.auth.assumable_iam_role:
auth_role = _common.AuthRole(
assumable_iam_role=pb2.auth.assumable_iam_role)
else:
auth_role = _common.AuthRole(
assumable_iam_role=pb2.auth.kubernetes_service_account)
return cls(
workflow_id=_identifier.Identifier.from_flyte_idl(pb2.workflow_id),
entity_metadata=LaunchPlanMetadata.from_flyte_idl(
pb2.entity_metadata),
default_inputs=_interface.ParameterMap.from_flyte_idl(
pb2.default_inputs),
fixed_inputs=_literals.LiteralMap.from_flyte_idl(pb2.fixed_inputs),
labels=_common.Labels.from_flyte_idl(pb2.labels),
annotations=_common.Annotations.from_flyte_idl(pb2.annotations),
auth_role=auth_role,
raw_output_data_config=_common.RawOutputDataConfig.from_flyte_idl(
pb2.raw_output_data_config),
)
def test_auth_role():
obj = _common.AuthRole(assumable_iam_role="rollie-pollie")
assert obj.assumable_iam_role == "rollie-pollie"
assert not obj.kubernetes_service_account
obj2 = _common.AuthRole.from_flyte_idl(obj.to_flyte_idl())
assert obj == obj2
obj = _common.AuthRole(kubernetes_service_account="service-account-name")
assert obj.kubernetes_service_account == "service-account-name"
assert not obj.assumable_iam_role
obj2 = _common.AuthRole.from_flyte_idl(obj.to_flyte_idl())
assert obj == obj2
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 test_auth_role_empty():
# This test is here to ensure we can serialize launch plans with an empty auth role.
# Auth roles are empty because they are filled in at registration time.
obj = _common.AuthRole()
x = obj.to_flyte_idl()
y = _common.AuthRole.from_flyte_idl(x)
assert y == obj
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 __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 launch(self, project, domain, name=None, inputs=None, notification_overrides=None, label_overrides=None,
annotation_overrides=None, auth_role=None):
"""
Executes the task as a single task execution and returns the identifier.
:param Text project:
:param Text domain:
:param Text name:
:param flytekit.models.literals.LiteralMap inputs: The inputs to pass
:param list[flytekit.models.common.Notification] notification_overrides: If specified, override the
notifications.
:param flytekit.models.common.Labels label_overrides:
:param flytekit.models.common.Annotations annotation_overrides:
:param flytekit.models.common.AuthRole auth_role:
:rtype: flytekit.models.execution.Execution
"""
disable_all = (notification_overrides == [])
if disable_all:
notification_overrides = None
else:
notification_overrides = _execution_models.NotificationList(
notification_overrides or []
)
disable_all = None
if not auth_role:
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):
_logging.warning("Using deprecated `role` from config. "
"Please update your config to use `assumable_iam_role` instead")
assumable_iam_role = _sdk_config.ROLE.get()
auth_role = _common_models.AuthRole(assumable_iam_role=assumable_iam_role,
kubernetes_service_account=kubernetes_service_account)
try:
# TODO(katrogan): Add handling to register the underlying task if it's not already.
client = _FlyteClientManager(_platform_config.URL.get(), insecure=_platform_config.INSECURE.get()).client
exec_id = client.create_execution(
project,
domain,
name,
_execution_models.ExecutionSpec(
self.sdk_task.id,
_execution_models.ExecutionMetadata(
_execution_models.ExecutionMetadata.ExecutionMode.MANUAL,
'sdk', # TODO: get principle
0 # TODO: Detect nesting
),
notifications=notification_overrides,
disable_all=disable_all,
labels=label_overrides,
annotations=annotation_overrides,
auth_role=auth_role,
),
inputs,
)
except _user_exceptions.FlyteEntityAlreadyExistsException:
exec_id = _identifier.WorkflowExecutionIdentifier(project, domain, name)
return client.get_execution(exec_id)
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,
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 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 auth_role(self):
"""
:rtype: flytekit.models.common.AuthRole
"""
fixed_auth = super(SdkLaunchPlan, self).auth_role
if fixed_auth is not None and (
fixed_auth.assumable_iam_role is not None or fixed_auth.kubernetes_service_account is not None
):
return fixed_auth
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):
_logging.warning(
"Using deprecated `role` from config. Please update your config to use `assumable_iam_role` instead"
)
assumable_iam_role = _sdk_config.ROLE.get()
return _common_models.AuthRole(
assumable_iam_role=assumable_iam_role, kubernetes_service_account=kubernetes_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 __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 launch_with_literals(
self,
project,
domain,
literal_inputs,
name=None,
notification_overrides=None,
label_overrides=None,
annotation_overrides=None,
):
"""
Launches a single task execution and returns the execution identifier.
:param Text project:
:param Text domain:
:param flytekit.models.literals.LiteralMap literal_inputs: Inputs to the execution.
:param Text name: [Optional] If specified, an execution will be created with this name. Note: the name must
be unique within the context of the project and domain.
:param list[flytekit.common.notifications.Notification] notification_overrides: [Optional] If specified, these
are the notifications that will be honored for this execution. An empty list signals to disable all
notifications.
:param flytekit.models.common.Labels label_overrides:
:param flytekit.models.common.Annotations annotation_overrides:
:rtype: flytekit.common.workflow_execution.SdkWorkflowExecution
"""
disable_all = notification_overrides == []
if disable_all:
notification_overrides = None
else:
notification_overrides = _admin_execution_models.NotificationList(
notification_overrides or [])
disable_all = None
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):
_logging.warning(
"Using deprecated `role` from config. "
"Please update your config to use `assumable_iam_role` instead"
)
assumable_iam_role = _sdk_config.ROLE.get()
auth_role = _common_model.AuthRole(
assumable_iam_role=assumable_iam_role,
kubernetes_service_account=kubernetes_service_account,
)
client = _flyte_engine.get_client()
try:
# TODO(katrogan): Add handling to register the underlying task if it's not already.
exec_id = client.create_execution(
project,
domain,
name,
_admin_execution_models.ExecutionSpec(
self.id,
_admin_execution_models.ExecutionMetadata(
_admin_execution_models.ExecutionMetadata.
ExecutionMode.MANUAL,
"sdk", # TODO: get principle
0, # TODO: Detect nesting
),
notifications=notification_overrides,
disable_all=disable_all,
labels=label_overrides,
annotations=annotation_overrides,
auth_role=auth_role,
),
literal_inputs,
)
except _user_exceptions.FlyteEntityAlreadyExistsException:
exec_id = _identifier.WorkflowExecutionIdentifier(
project, domain, name)
execution = client.get_execution(exec_id)
return _workflow_execution.SdkWorkflowExecution.promote_from_model(
execution)
