def test_hydrate_registration_parameters__workflow_nothing_set(): workflow = _workflow_pb2.WorkflowSpec( template=_core_workflow_pb2.WorkflowTemplate( id=_identifier_pb2.Identifier( resource_type=_identifier_pb2.WORKFLOW, name="name", ), nodes=[ _core_workflow_pb2.Node( id="foo", task_node=_core_workflow_pb2.TaskNode( reference_id=_identifier_pb2.Identifier( resource_type=_identifier_pb2.TASK, name="task1")), ) ], )) identifier, entity = hydrate_registration_parameters( _identifier_pb2.WORKFLOW, "project", "domain", "12345", workflow) assert (identifier == _identifier_pb2.Identifier( resource_type=_identifier_pb2.WORKFLOW, project="project", domain="domain", name="name", version="12345", ) == entity.template.id) assert len(workflow.template.nodes) == 1 assert workflow.template.nodes[ 0].task_node.reference_id == _identifier_pb2.Identifier( resource_type=_identifier_pb2.TASK, project="project", domain="domain", name="task1", version="12345", )
def test_hydrate_registration_parameters__launch_plan_nothing_set(): launch_plan = _launch_plan_pb2.LaunchPlan( id=_identifier_pb2.Identifier( resource_type=_identifier_pb2.LAUNCH_PLAN, name="lp_name", ), spec=_launch_plan_pb2.LaunchPlanSpec( workflow_id=_identifier_pb2.Identifier( resource_type=_identifier_pb2.WORKFLOW, name="workflow_name", )), ) identifier, entity = hydrate_registration_parameters( _identifier_pb2.LAUNCH_PLAN, "project", "domain", "12345", launch_plan, ) assert identifier == _identifier_pb2.Identifier( resource_type=_identifier_pb2.LAUNCH_PLAN, project="project", domain="domain", name="lp_name", version="12345", ) assert entity.spec.workflow_id == _identifier_pb2.Identifier( resource_type=_identifier_pb2.WORKFLOW, project="project", domain="domain", name="workflow_name", version="12345", )
def test_hydrate_registration_parameters__launch_plan_already_set(): launch_plan = _launch_plan_pb2.LaunchPlanSpec( workflow_id=_identifier_pb2.Identifier( resource_type=_identifier_pb2.WORKFLOW, project="project2", domain="domain2", name="workflow_name", version="abc", ) ) identifier, entity = hydrate_registration_parameters( _identifier_pb2.Identifier( resource_type=_identifier_pb2.LAUNCH_PLAN, project="project2", domain="domain2", name="workflow_name", version="abc", ), "project", "domain", "12345", launch_plan, ) assert identifier == _identifier_pb2.Identifier( resource_type=_identifier_pb2.LAUNCH_PLAN, project="project2", domain="domain2", name="workflow_name", version="abc", ) assert entity.workflow_id == launch_plan.workflow_id
def register( registrable_entities: typing.List[RegistrableEntity], project: str, domain: str, version: str, client: SynchronousFlyteClient, ): # The incoming registrable entities are already in base protobuf form, not model form, so we use the # raw client's methods instead of the friendly client's methods by calling super for admin_entity in registrable_entities: try: if isinstance(admin_entity, _idl_admin_TaskSpec): ident, task_spec = hydrate_registration_parameters( identifier_pb2.TASK, project, domain, version, admin_entity) logger.debug(f"Creating task {ident}") super(SynchronousFlyteClient, client).create_task( TaskCreateRequest(id=ident, spec=task_spec)) elif isinstance(admin_entity, _idl_admin_WorkflowSpec): ident, wf_spec = hydrate_registration_parameters( identifier_pb2.WORKFLOW, project, domain, version, admin_entity) logger.debug(f"Creating workflow {ident}") super(SynchronousFlyteClient, client).create_workflow( WorkflowCreateRequest(id=ident, spec=wf_spec)) elif isinstance(admin_entity, _idl_admin_LaunchPlan): ident, admin_lp = hydrate_registration_parameters( identifier_pb2.LAUNCH_PLAN, project, domain, version, admin_entity) logger.debug(f"Creating launch plan {ident}") super(SynchronousFlyteClient, client).create_launch_plan( LaunchPlanCreateRequest(id=ident, spec=admin_lp.spec)) else: raise AssertionError( f"Unknown entity of type {type(admin_entity)}") except FlyteEntityAlreadyExistsException: logger.info(f"{admin_entity} already exists") except Exception as e: logger.info( f"Failed to register entity {admin_entity} with error {e}") raise e
def test_hydrate_registration_parameters__task_nothing_set(): task = _task_pb2.TaskSpec( template=_core_task_pb2.TaskTemplate( id=_identifier_pb2.Identifier(resource_type=_identifier_pb2.TASK, name="name",), ) ) identifier, entity = hydrate_registration_parameters(task.template.id, "project", "domain", "12345", task) assert ( identifier == _identifier_pb2.Identifier( resource_type=_identifier_pb2.TASK, project="project", domain="domain", name="name", version="12345", ) == entity.template.id )
def test_hydrate_registration_parameters__workflow_already_set(): workflow = _workflow_pb2.WorkflowSpec( template=_core_workflow_pb2.WorkflowTemplate( id=_identifier_pb2.Identifier( resource_type=_identifier_pb2.WORKFLOW, project="project2", domain="domain2", name="name", version="abc", ), )) identifier, entity = hydrate_registration_parameters( _identifier_pb2.WORKFLOW, "project", "domain", "12345", workflow) assert (identifier == _identifier_pb2.Identifier( resource_type=_identifier_pb2.WORKFLOW, project="project2", domain="domain2", name="name", version="abc", ) == entity.template.id)
def test_hydrate_registration_parameters__subworkflows(): workflow_template = _core_workflow_pb2.WorkflowTemplate() workflow_template.id.CopyFrom( _identifier_pb2.Identifier(resource_type=_identifier_pb2.WORKFLOW, name="workflow")) sub_workflow_template = _core_workflow_pb2.WorkflowTemplate() sub_workflow_template.id.CopyFrom( _identifier_pb2.Identifier(resource_type=_identifier_pb2.WORKFLOW, name="subworkflow")) sub_workflow_template.nodes.append( _core_workflow_pb2.Node( id="task_node", task_node=_core_workflow_pb2.TaskNode( reference_id=_identifier_pb2.Identifier( resource_type=_identifier_pb2.TASK)), )) workflow_spec = _workflow_pb2.WorkflowSpec(template=workflow_template) workflow_spec.sub_workflows.append(sub_workflow_template) identifier, entity = hydrate_registration_parameters( _identifier_pb2.WORKFLOW, "project", "domain", "12345", workflow_spec) assert (identifier == _identifier_pb2.Identifier( resource_type=_identifier_pb2.WORKFLOW, project="project", domain="domain", name="workflow", version="12345", ) == entity.template.id) assert entity.sub_workflows[0].id == _identifier_pb2.Identifier( resource_type=_identifier_pb2.WORKFLOW, project="project", domain="domain", name="subworkflow", version="12345", )