def testRegisterWrongVar(self):
# Have a var not exist in the function and assert there is
# an error
create_component("test_component", "test_description", "shreya")
@register(
component_name="test_component",
input_vars=["foory"], # Not valid
output_vars=["bar"],
)
def test_func():
x = 0
foo = x + 1
bar = x + 2
@register(
component_name="test_component",
input_vars=["foo"],
output_vars=["barry"], # Not valid
)
def test_func2():
x = 0
foo = x + 1
bar = x + 2
with self.assertRaises(ValueError):
test_func()
with self.assertRaises(ValueError):
test_func2()
def testLogEmptyComponentRun(self):
# Create component then log a run of it
create_component("test_component", "test_description", "shreya")
# Create a ComponentRun
cr = ComponentRun("test_component")
with self.assertRaises(RuntimeError):
log_component_run(cr)
def testRegisterNoReturn(self):
# Create component then log a run of it
create_component("test_component", "test_description", "shreya")
@register(component_name="test_component",
input_vars=["foo"],
output_vars=["bar"])
def test_func():
x = 0
foo = x + 1
bar = x + 2
test_func()
def testLogBasicComponentRun(self):
# Create component then log a run of it
create_component("test_component", "test_description", "shreya")
# Create a ComponentRun
cr = ComponentRun(component_name="test_component")
cr.set_start_timestamp()
cr.code_snapshot = b"def main(): return"
cr.add_inputs(["duplicate_input", "duplicate_input"])
cr.add_outputs(["duplicate_output", "duplicate_output"])
cr.set_end_timestamp()
# Log component run
log_component_run(cr)
def testRegister(self):
# Create component then log a run of it
create_component("test_component", "test_description", "shreya")
@register(
component_name="test_component",
input_vars=["foo"],
output_vars=["bar"],
)
def test_func():
foo = "foo"
bar = "bar"
return
test_func()
def testLogKVComponentRun(self):
# Tests implementation of values in iopointer
create_component(
name="valtest",
description="Tests implementation of values in iopointer.",
owner="me",
)
iop1 = ["this", "is", "the", "first"]
iop2 = ["this", "is", "the", "second"]
# Create iopointers and CR
iop1 = IOPointer(name="iop1", value=iop1)
iop2 = IOPointer(name="iop2", value=iop2)
cr = ComponentRun("valtest")
cr.set_start_timestamp()
cr.set_end_timestamp()
cr.add_input(iop1)
cr.add_output(iop2)
log_component_run(cr)
input_vars=["filename", "model_path"],
output_vars=["output_id"],
)
def inference(filename: str, model_path: str) -> str:
# Load model and data
# Run some inference
output_id = "".join(
random.choice(string.ascii_lowercase) for i in range(10))
return output_id
if __name__ == "__main__":
# Create components
create_component(
name="ingest",
description="Example of ingesting data from some client.",
owner="data_engineer",
tags=["example"],
)
create_component(
name="clean",
description="Example of cleaning data.",
owner="data_engineer",
tags=["example"],
)
create_component(
name="featuregen",
description="Example of generating features.",
owner="ml_engineer",
tags=["example"],
)
create_component(
component_name="research_model_development",
input_vars=["lr", "num_epochs", "hidden_size"],
output_vars=["model_metric"],
)
def train_and_evaluate_model(lr, num_epochs, hidden_size):
# Grab some random accuracy
model_metric = f"{random.uniform(0, 1)}_accuracy"
return model_metric
if __name__ == "__main__":
# Create component
create_component(
name="research_model_development",
description="Example of training a model for research.",
owner="neurips_queen",
tags=["example"],
)
lr_candidates = [0.001, 0.01, 0.1]
num_epochs_candidates = [10, 50, 100]
hidden_size_candidates = [64, 128, 256]
candidates = itertools.product(
lr_candidates, num_epochs_candidates, hidden_size_candidates
)
for lr, num_epochs, hidden_size in candidates:
train_and_evaluate_model(
f"lr_{lr}", f"num_epochs_{num_epochs}", f"hidden_size_{hidden_size}"
)
from mltrace import create_component, register
import random
import string
_identifier = "".join(random.choice(string.ascii_lowercase) for i in range(10))
@register(component_name="tiny",
input_vars=["inp_str"],
output_vars=["out_str"])
def increment(inp: int) -> int:
inp_str = f"{_identifier}_{str(inp)}"
out_str = f"{_identifier}_{str(inp + 1)}"
return inp + 1
if __name__ == "__main__":
# Create component
create_component(
name="tiny",
description="Example of a tiny component",
owner="groot",
tags=["example"],
)
# Run the tiny function with some fake inputs and outputs
i = 0
while i < 10:
i = increment(i)
