def func_m(
a: int,
b: int,
num_iterations: Parameter[int] = 'abc'
) -> OutputDict(c=float):
del num_iterations
return {'c': float(a + b)}
def upload_model(project_id: Parameter[str], display_name: Parameter[str],
serving_container: Parameter[str], region: Parameter[str],
model: InputArtifact[Model]) -> OutputDict(model_name=str):
"""Uploads model artifacts to AI Platform Models."""
api_endpoint = f'{region}-aiplatform.googleapis.com'
parent = f'projects/{project_id}/locations/{region}'
client_options = {'api_endpoint': api_endpoint}
client = aiplatform.gapic.ModelServiceClient(client_options=client_options)
artifact_uri = '{}/serving_model_dir'.format(model.uri)
model_spec = {
'display_name': display_name,
'metadata_schema_uri': "",
'artifact_uri': artifact_uri,
'container_spec': {
'image_uri': serving_container,
'command': [],
'args': []
}
}
response = client.upload_model(parent=parent, model=model_spec)
logging.info('Uploading model {}. Operation ID: {}'.format(
model, response.operation.name))
upload_model_response = response.result()
logging.info('Upload completed.')
return {'model_name': upload_model_response}
def func_j(
a: int,
b: int,
examples: InputArtifact[standard_artifacts.Examples] = 123
) -> OutputDict(c=float):
del examples
return {'c': float(a + b)}
def func_k(
a: int,
b: int,
model: OutputArtifact[standard_artifacts.Model] = None
) -> OutputDict(c=float):
del model
return {'c': float(a + b)}
def func_a() -> OutputDict(precision=float,
recall=float,
message=str,
serialized_value=bytes,
optional_label=Optional[str],
optional_metric=Optional[float]):
return {
'precision': 0.9,
'recall': 0.8,
'message': 'foo',
'serialized_value': b'bar',
'optional_label': None,
'optional_metric': 1.0,
}
def func_a(
examples: InputArtifact[standard_artifacts.Examples],
model: OutputArtifact[standard_artifacts.Model],
schema: InputArtifact[standard_artifacts.Schema],
statistics: OutputArtifact[standard_artifacts.ExampleStatistics],
num_steps: Parameter[int]
) -> OutputDict(
precision=float, recall=float, message=Text, serialized_value=bytes):
del examples, model, schema, statistics, num_steps
return {
'precision': 0.9,
'recall': 0.8,
'message': 'foo',
'serialized_value': b'bar'
}
def DummyTrainComponent(
training_data: InputArtifact[DummyDataset],
model: OutputArtifact[DummyModel],
num_iterations: Parameter[int] = 10) -> OutputDict(
loss=float, accuracy=float):
"""Simple fake trainer component."""
records = training_data.read()
model_obj, loss, accuracy = train_dummy_model(records, num_iterations)
model.write(model_obj)
LocalDagRunnerTest.RAN_COMPONENTS.append('Train')
return {
'loss': loss,
'accuracy': accuracy,
}
def func_h(a: int,
b: int) -> OutputDict(c=standard_artifacts.Examples):
return {'c': float(a + b)}
def func_g(a: int, b: int) -> OutputDict(c='whatever'):
return {'c': float(a + b)}
def func_f(a: int, b: Dict[int, int]) -> OutputDict(c=float):
return {'c': float(a + b)}
def func_e(
a: int, unused_b: standard_artifacts.Examples
) -> OutputDict(c=float):
return {'c': float(a)}
def _simple_component(a: int, b: int, c: str,
d: bytes) -> OutputDict(e=float, f=float):
del c, d
return {'e': float(a + b), 'f': float(a * b)}
def _injector_1(
foo: Parameter[int],
bar: Parameter[str]) -> OutputDict(a=int, b=int, c=str, d=bytes):
assert foo == 9
assert bar == 'secret'
return {'a': 10, 'b': 22, 'c': 'unicode', 'd': b'bytes'}
def _simple_component(
a: int, b: int, c: str, d: bytes
) -> OutputDict(e=float, f=float, g=Optional[str], h=Optional[str]):
del c, d
return {'e': float(a + b), 'f': float(a * b), 'g': 'OK', 'h': None}
def func_i(
a: int, b: int
) -> OutputDict(c=OutputArtifact[standard_artifacts.Examples]):
return {'c': float(a + b)}
def _injector_2(
examples: OutputArtifact[standard_artifacts.Examples]
) -> OutputDict(a=int, b=float, c=str, d=bytes, e=str):
fileio.makedirs(examples.uri)
return {'a': 1, 'b': 2.0, 'c': '3', 'd': b'4', 'e': 'passed'}
def func_c(a: int, b: int, **unused_kwargs) -> OutputDict(c=float):
return {'c': float(a + b)}
def func_l(a: int,
b: int,
num_iterations: int = 'abc') -> OutputDict(c=float):
del num_iterations
return {'c': float(a + b)}
def func_d(a: int, b) -> OutputDict(c=float):
return {'c': float(a + b)}
def func_a(a: int, b: int, unused_c: Text, unused_d: bytes,
unused_e: Parameter[float]) -> OutputDict(c=float):
return {'c': float(a + b)}
def _injector_2(
examples: OutputArtifact[standard_artifacts.Examples]
) -> OutputDict(a=int, b=float, c=Text, d=bytes, e=Text):
del examples
return {'a': 1, 'b': 2.0, 'c': '3', 'd': b'4', 'e': 'passed'}
