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",
)