def test_construct_no_destination_custom_executor(self):
pusher = component.Pusher(
model_export=self.model_export,
model_blessing=self.model_blessing,
executor_class=self._MyCustomPusherExecutor,
)
self.assertEqual('ModelPushPath', pusher.outputs.model_push.type_name)
def testEnableCache(self):
pusher_1 = component.Pusher(
model=self.model,
model_blessing=self.model_blessing,
push_destination=pusher_pb2.PushDestination(
filesystem=pusher_pb2.PushDestination.Filesystem(
base_directory='push_destination')))
self.assertEqual(None, pusher_1.enable_cache)
pusher_2 = component.Pusher(
model=self.model,
model_blessing=self.model_blessing,
push_destination=pusher_pb2.PushDestination(
filesystem=pusher_pb2.PushDestination.Filesystem(
base_directory='push_destination')),
enable_cache=True)
self.assertEqual(True, pusher_2.enable_cache)
def testConstruct(self):
pusher = component.Pusher(model=self._model,
model_blessing=self._model_blessing,
push_destination=self._push_destination)
self.assertEqual(
standard_artifacts.PushedModel.TYPE_NAME, pusher.outputs[
standard_component_specs.PUSHED_MODEL_KEY].type_name)
def test_construct(self):
pusher = component.Pusher(
model_export=self.model_export,
model_blessing=self.model_blessing,
push_destination=pusher_pb2.PushDestination(
filesystem=pusher_pb2.PushDestination.Filesystem(
base_directory='push_destination')))
self.assertEqual('ModelPushPath', pusher.outputs.model_push.type_name)
def testConstructNoDestinationCustomExecutor(self):
pusher = component.Pusher(
model_export=self.model_export,
model_blessing=self.model_blessing,
custom_executor_spec=executor_spec.ExecutorClassSpec(
self._MyCustomPusherExecutor),
)
self.assertEqual('ModelPushPath', pusher.outputs['model_push'].type_name)
def testConstructNoDestinationCustomExecutor(self):
pusher = component.Pusher(
model=self.model,
model_blessing=self.model_blessing,
custom_executor_spec=executor_spec.ExecutorClassSpec(
self._MyCustomPusherExecutor),
)
self.assertEqual(standard_artifacts.PushedModel.TYPE_NAME,
pusher.outputs['pushed_model'].type_name)
def testConstruct(self):
pusher = component.Pusher(
model=self.model,
model_blessing=self.model_blessing,
push_destination=pusher_pb2.PushDestination(
filesystem=pusher_pb2.PushDestination.Filesystem(
base_directory='push_destination')))
self.assertEqual(standard_artifacts.PushedModel.TYPE_NAME,
pusher.outputs['pushed_model'].type_name)
def testConstruct_NoModelAndNoInfraBlessing_Fails(self):
with self.assertRaisesRegex(
ValueError,
('Either one of model or infra_blessing channel should be given')):
component.Pusher(
# model=self._model, # No model.
model_blessing=self._model_blessing,
# infra_blessing=self._infra_blessing, # No infra_blessing.
push_destination=self._push_destination)
def testConstruct_InfraBlessingReplacesModel(self):
pusher = component.Pusher(
# model=self._model, # No model.
model_blessing=self._model_blessing,
infra_blessing=self._infra_blessing,
push_destination=self._push_destination)
self.assertCountEqual(pusher.inputs.keys(),
['model_blessing', 'infra_blessing'])
def testConstructWithParameter(self):
push_dir = data_types.RuntimeParameter(name='push-dir', ptype=Text)
pusher = component.Pusher(
model=self.model,
model_blessing=self.model_blessing,
push_destination={'filesystem': {
'base_directory': push_dir
}})
self.assertEqual(standard_artifacts.PushedModel.TYPE_NAME,
pusher.outputs['pushed_model'].type_name)
def test_construct(self):
model_export = types.TfxArtifact(type_name='ModelExportPath')
model_blessing = types.TfxArtifact(type_name='ModelBlessingPath')
pusher = component.Pusher(
model_export=channel.as_channel([model_export]),
model_blessing=channel.as_channel([model_blessing]),
push_destination=pusher_pb2.PushDestination(
filesystem=pusher_pb2.PushDestination.Filesystem(
base_directory='push_destination')))
self.assertEqual('ModelPushPath', pusher.outputs.model_push.type_name)
def testConstructWithParameter(self):
push_dir = data_types.RuntimeParameter(name='push-dir', ptype=Text)
pusher = component.Pusher(
model=self.model,
model_blessing=self.model_blessing,
push_destination={'filesystem': {
'base_directory': push_dir
}})
self.assertEqual('ModelPushPath',
pusher.outputs['model_push'].type_name)
def __init__(self, model_export: str, model_blessing: str,
serving_directory: str):
push_destination = pusher_pb2.PushDestination(
filesystem=pusher_pb2.PushDestination.Filesystem(
base_directory=serving_directory))
component = pusher_component.Pusher(
model_export=channel.Channel('ModelExportPath'),
model_blessing=channel.Channel('ModelBlessingPath'),
push_destination=push_destination)
super().__init__(component, {
"model_export": model_export,
"model_blessing": model_blessing,
})
def testConstructNoDestination(self):
with self.assertRaises(ValueError):
_ = component.Pusher(
model=self.model,
model_blessing=self.model_blessing,
)
def test_construct_no_destination(self):
with self.assertRaises(ValueError):
_ = component.Pusher(
model_export=self.model_export,
model_blessing=self.model_blessing,
)
