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)