From 6b28b63c82c35df230def7ac8384a5bdb277a72b Mon Sep 17 00:00:00 2001 From: tfx-team Date: Thu, 27 Jun 2024 14:49:08 -0700 Subject: [PATCH] add guri to importer artifacts PiperOrigin-RevId: 647454281 --- .../core/task_schedulers/importer_task_scheduler.py | 2 ++ .../core/task_schedulers/importer_task_scheduler_test.py | 3 +++ 2 files changed, 5 insertions(+) diff --git a/tfx/orchestration/experimental/core/task_schedulers/importer_task_scheduler.py b/tfx/orchestration/experimental/core/task_schedulers/importer_task_scheduler.py index 5e47a12c084..4a75392165a 100644 --- a/tfx/orchestration/experimental/core/task_schedulers/importer_task_scheduler.py +++ b/tfx/orchestration/experimental/core/task_schedulers/importer_task_scheduler.py @@ -23,6 +23,7 @@ from tfx.utils import status as status_lib + class ImporterTaskScheduler(task_scheduler.TaskScheduler[task_lib.ExecNodeTask] ): """A task scheduler for Importer system node.""" @@ -49,6 +50,7 @@ def schedule(self) -> task_scheduler.TaskSchedulerResult: output_key=output_key, ) + return task_scheduler.TaskSchedulerResult( status=status_lib.Status(code=status_lib.Code.OK), output=task_scheduler.ImporterNodeOutput( diff --git a/tfx/orchestration/experimental/core/task_schedulers/importer_task_scheduler_test.py b/tfx/orchestration/experimental/core/task_schedulers/importer_task_scheduler_test.py index 0fe8514d8ed..6824b45fab6 100644 --- a/tfx/orchestration/experimental/core/task_schedulers/importer_task_scheduler_test.py +++ b/tfx/orchestration/experimental/core/task_schedulers/importer_task_scheduler_test.py @@ -31,6 +31,8 @@ from tfx.utils import status as status_lib + + class ImporterTaskSchedulerTest(test_utils.TfxTest): def setUp(self): @@ -80,6 +82,7 @@ def _make_pipeline(self, pipeline_root, pipeline_run_id): return pipeline def test_importer_task_scheduler(self): + with self._mlmd_connection as m: ts_result = importer_task_scheduler.ImporterTaskScheduler( mlmd_handle=m, pipeline=self._pipeline,