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)