Skip to content

Commit

Permalink
no-op
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 660551866
  • Loading branch information
tfx-copybara committed Aug 7, 2024
1 parent 236ac38 commit 3a8015b
Show file tree
Hide file tree
Showing 3 changed files with 44 additions and 9 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -490,6 +490,8 @@ def _generate_tasks_for_node(
execution_type=node.node_info.type,
contexts=resolved_info.contexts,
input_and_params=unprocessed_inputs,
pipeline=self._pipeline,
node_id=node.node_info.id,
)

for execution in executions:
Expand Down
2 changes: 2 additions & 0 deletions tfx/orchestration/experimental/core/sync_pipeline_task_gen.py
Original file line number Diff line number Diff line change
Expand Up @@ -564,6 +564,8 @@ def _generate_tasks_from_resolved_inputs(
execution_type=node.node_info.type,
contexts=resolved_info.contexts,
input_and_params=resolved_info.input_and_params,
pipeline=self._pipeline,
node_id=node.node_info.id,
)

result.extend(
Expand Down
49 changes: 40 additions & 9 deletions tfx/orchestration/experimental/core/task_gen_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@
from tfx.orchestration import metadata
from tfx.orchestration import node_proto_view
from tfx.orchestration.experimental.core import constants
from tfx.orchestration.experimental.core import env
from tfx.orchestration.experimental.core import mlmd_state
from tfx.orchestration.experimental.core import task as task_lib
from tfx.orchestration import mlmd_connection_manager as mlmd_cm
Expand Down Expand Up @@ -548,21 +549,35 @@ def register_executions_from_existing_executions(
contexts = metadata_handle.store.get_contexts_by_execution(
existing_executions[0].id
)
return execution_lib.put_executions(
executions = execution_lib.put_executions(
metadata_handle,
new_executions,
contexts,
input_artifacts_maps=input_artifacts,
)

pipeline_asset = metadata_handle.store.pipeline_asset
if pipeline_asset:
env.get_env().create_pipeline_run_node_executions(
pipeline_asset.owner,
pipeline_asset.name,
pipeline,
node.node_info.id,
executions,
)

return executions


def register_executions(
metadata_handle: metadata.Metadata,
execution_type: metadata_store_pb2.ExecutionType,
contexts: Sequence[metadata_store_pb2.Context],
input_and_params: Sequence[InputAndParam],
pipeline: Optional[pipeline_pb2.Pipeline] = None,
node_id: Optional[str] = None,
) -> Sequence[metadata_store_pb2.Execution]:
"""Registers multiple executions in MLMD.
"""Registers multiple executions in MLMD and Tflex backends.
Along with the execution:
- the input artifacts will be linked to the executions.
Expand All @@ -575,6 +590,8 @@ def register_executions(
input_and_params: A list of InputAndParams, which includes input_dicts
(dictionaries of artifacts. One execution will be registered for each of
the input_dict) and corresponding exec_properties.
pipeline: Optional. The pipeline proto.
node_id: Optional. The node id of the executions to be registered.
Returns:
A list of MLMD executions that are registered in MLMD, with id populated.
Expand Down Expand Up @@ -603,21 +620,35 @@ def register_executions(
executions.append(execution)

if len(executions) == 1:
return [
new_executions = [
execution_lib.put_execution(
metadata_handle,
executions[0],
contexts,
input_artifacts=input_and_params[0].input_artifacts,
)
]
else:
new_executions = execution_lib.put_executions(
metadata_handle,
executions,
contexts,
[
input_and_param.input_artifacts
for input_and_param in input_and_params
],
)

return execution_lib.put_executions(
metadata_handle,
executions,
contexts,
[input_and_param.input_artifacts for input_and_param in input_and_params],
)
pipeline_asset = metadata_handle.store.pipeline_asset
if pipeline_asset and pipeline and node_id:
env.get_env().create_pipeline_run_node_executions(
pipeline_asset.owner,
pipeline_asset.name,
pipeline,
node_id,
new_executions,
)
return new_executions


def update_external_artifact_type(
Expand Down

0 comments on commit 3a8015b

Please sign in to comment.