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_metadata_schedule(): s = schedule.Schedule('asdf', '1 3 4 5 6 7') obj = launch_plan.LaunchPlanMetadata(schedule=s, notifications=[]) assert obj.schedule == s obj2 = launch_plan.LaunchPlanMetadata.from_flyte_idl(obj.to_flyte_idl()) assert obj == obj2 assert obj2.schedule == s
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 test_schedule_cron_schedule(offset): cs = _schedule.Schedule.CronSchedule("days", offset) obj = _schedule.Schedule(cron_schedule=cs, kickoff_time_input_arg="fdsa") assert obj.cron_schedule.schedule == "days" assert obj.schedule_expression == cs assert obj.rate is None assert obj.cron_expression is None obj2 = _schedule.Schedule.from_flyte_idl(obj.to_flyte_idl()) assert obj == obj2 assert obj2.schedule_expression == cs assert obj.rate is None assert obj.cron_expression is None
def test_schedule(): obj = _schedule.Schedule(kickoff_time_input_arg="fdsa", cron_expression="1 2 3 4 5 6") assert obj.rate is None assert obj.cron_expression == "1 2 3 4 5 6" assert obj.schedule_expression == "1 2 3 4 5 6" assert obj.kickoff_time_input_arg == "fdsa" obj2 = _schedule.Schedule.from_flyte_idl(obj.to_flyte_idl()) assert obj == obj2 assert obj2.rate is None assert obj2.cron_expression == "1 2 3 4 5 6" assert obj2.schedule_expression == "1 2 3 4 5 6" assert obj2.kickoff_time_input_arg == "fdsa"
def test_schedule_with_lp(): @task def double(a: int) -> int: return a * 2 @workflow def quadruple(a: int) -> int: b = double(a=a) c = double(a=b) return c lp = LaunchPlan.create( "schedule_test", quadruple, schedule=FixedRate(_datetime.timedelta(hours=12), "kickoff_input"), ) assert lp.schedule == _schedule_models.Schedule( "kickoff_input", rate=_schedule_models.Schedule.FixedRate(12, _schedule_models.Schedule.FixedRateUnit.HOUR) )
def test_schedule_fixed_rate(): fr = _schedule.Schedule.FixedRate(10, _schedule.Schedule.FixedRateUnit.MINUTE) obj = _schedule.Schedule(kickoff_time_input_arg="fdsa", rate=fr) assert obj.rate is not None assert obj.cron_expression is None assert obj.kickoff_time_input_arg == "fdsa" assert obj.rate == fr assert obj.schedule_expression == fr obj2 = _schedule.Schedule.from_flyte_idl(obj.to_flyte_idl()) assert obj == obj2 assert obj2.rate is not None assert obj2.cron_expression is None assert obj2.kickoff_time_input_arg == "fdsa" assert obj2.rate == fr assert obj2.schedule_expression == fr
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 __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