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