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 __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 test_execution_annotation_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")) annotations = _common_models.Annotations({"my": "annotation"}) engine.FlyteLaunchPlan(m).execute('xp', 'xd', 'xn', literals.LiteralMap({}), notification_overrides=[], annotation_overrides=annotations) 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, annotations=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, 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_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_annotations(): 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', annotations=_common_models.Annotations({"my": "annotation"})) assert lp.annotations.values == {"my": "annotation"}
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_annotations(): obj = _common.Annotations({"my": "annotation"}) assert obj.values == {"my": "annotation"} obj2 = _common.Annotations.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