def test_unused_node(default_model_storage: ModelStorage): graph_schema = GraphSchema({ "provide": SchemaNode( needs={}, uses=ProvideX, fn="provide", constructor_name="create", config={}, is_target=True, ), # This node will not fail as it will be pruned because it is not a target # or a target's ancestor. "assert_false": SchemaNode( needs={"i": "input"}, uses=AssertComponent, fn="run_assert", constructor_name="create", config={"value_to_assert": "some_value"}, ), }) runner = DaskGraphRunner( graph_schema=graph_schema, model_storage=default_model_storage, execution_context=ExecutionContext(graph_schema=graph_schema, model_id="1"), ) results = runner.run(inputs={"input": "some_other_value"}) assert results == {"provide": 1}
def test_target_override(eager: bool, default_model_storage: ModelStorage): graph_schema = GraphSchema( { "add": SchemaNode( needs={"i1": "first_input", "i2": "second_input"}, uses=AddInputs, fn="add", constructor_name="create", config={}, eager=eager, ), "subtract_2": SchemaNode( needs={"i": "add"}, uses=SubtractByX, fn="subtract_x", constructor_name="create", config={"x": 3}, eager=eager, is_target=True, ), } ) execution_context = ExecutionContext(graph_schema=graph_schema, model_id="1") runner = DaskGraphRunner( graph_schema=graph_schema, model_storage=default_model_storage, execution_context=execution_context, ) results = runner.run(inputs={"first_input": 3, "second_input": 4}, targets=["add"]) assert results == {"add": 7}
def test_loop(default_model_storage: ModelStorage): graph_schema = GraphSchema({ "subtract_a": SchemaNode( needs={"i": "subtract_b"}, uses=SubtractByX, fn="subtract_x", constructor_name="create", config={}, is_target=False, ), "subtract_b": SchemaNode( needs={"i": "subtract_a"}, uses=SubtractByX, fn="subtract_x", constructor_name="create", config={}, is_target=True, ), }) runner = DaskGraphRunner( graph_schema=graph_schema, model_storage=default_model_storage, execution_context=ExecutionContext(graph_schema=graph_schema, model_id="1"), ) with pytest.raises(GraphRunError): runner.run()
def test_unused_node(default_model_storage: ModelStorage): graph_schema = GraphSchema({ "provide": SchemaNode( needs={}, uses=ProvideX, fn="provide", constructor_name="create", config={}, is_target=True, ), "provide_2": SchemaNode( # This will not output needs={}, uses=ProvideX, fn="provide", constructor_name="create", config={}, ), }) runner = DaskGraphRunner( graph_schema=graph_schema, model_storage=default_model_storage, execution_context=ExecutionContext(graph_schema=graph_schema, model_id="1"), ) results = runner.run() assert results == {"provide": 1}
def test_empty_schema(default_model_storage: ModelStorage): empty_schema = GraphSchema({}) runner = DaskGraphRunner( graph_schema=empty_schema, model_storage=default_model_storage, execution_context=ExecutionContext(graph_schema=empty_schema, model_id="1"), ) results = runner.run() assert not results
def test_no_target(default_model_storage: ModelStorage): graph_schema = GraphSchema({ "provide": SchemaNode( needs={}, uses=ProvideX, fn="provide", constructor_name="create", config={}, ) }) runner = DaskGraphRunner( graph_schema=graph_schema, model_storage=default_model_storage, execution_context=ExecutionContext(graph_schema=graph_schema, model_id="1"), ) results = runner.run() assert not results
def test_input_value_is_node_name(default_model_storage: ModelStorage): graph_schema = GraphSchema({ "provide": SchemaNode( needs={}, uses=ProvideX, fn="provide", constructor_name="create", config={}, is_target=True, ) }) runner = DaskGraphRunner( graph_schema=graph_schema, model_storage=default_model_storage, execution_context=ExecutionContext(graph_schema=graph_schema, model_id="1"), ) with pytest.raises(GraphRunError): runner.run(inputs={"input": "provide"})
def test_can_use_alternate_constructor(default_model_storage: ModelStorage): graph_schema = GraphSchema({ "provide": SchemaNode( needs={}, uses=ProvideX, fn="provide", constructor_name="create_with_2", config={}, is_target=True, ) }) runner = DaskGraphRunner( graph_schema=graph_schema, model_storage=default_model_storage, execution_context=ExecutionContext(graph_schema=graph_schema, model_id="1"), ) results = runner.run() assert results["provide"] == 2
def test_non_eager_can_use_inputs_for_constructor(default_model_storage: ModelStorage): graph_schema = GraphSchema( { "provide": SchemaNode( needs={"x": "input"}, uses=ProvideX, fn="provide", constructor_name="create", config={}, eager=False, is_target=True, ), } ) runner = DaskGraphRunner( graph_schema=graph_schema, model_storage=default_model_storage, execution_context=ExecutionContext(graph_schema=graph_schema, model_id="1"), ) results = runner.run(inputs={"input": 5}) assert results["provide"] == 5
def test_execution_context(default_model_storage: ModelStorage): graph_schema = GraphSchema({ "execution_context_aware": SchemaNode( needs={}, uses=ExecutionContextAware, fn="get_execution_context", constructor_name="create", config={}, is_target=True, ) }) context = ExecutionContext(graph_schema=graph_schema, model_id="some_id") runner = DaskGraphRunner( graph_schema=graph_schema, model_storage=default_model_storage, execution_context=context, ) context.model_id = "a_new_id" result = runner.run()["execution_context_aware"] assert result.model_id == "some_id" assert result.node_name == "execution_context_aware"
def test_default_config(x: Optional[int], output: int, default_model_storage: ModelStorage): graph_schema = GraphSchema({ "subtract": SchemaNode( needs={"i": "input"}, uses=SubtractByX, fn="subtract_x", constructor_name="create", config={"x": x} if x else {}, is_target=True, ) }) runner = DaskGraphRunner( graph_schema=graph_schema, model_storage=default_model_storage, execution_context=ExecutionContext(graph_schema=graph_schema, model_id="1"), ) results = runner.run(inputs={"input": 5}) assert results["subtract"] == output
def test_loading_from_previous_node(default_model_storage: ModelStorage): test_value_for_sub_directory = {"test": "test value sub dir"} test_value = {"test dir": "test value dir"} graph_schema = GraphSchema({ "train": SchemaNode( needs={}, uses=PersistableTestComponent, fn="train", constructor_name="create", config={ "test_value": test_value, "test_value_for_sub_directory": test_value_for_sub_directory, }, ), "load": SchemaNode( needs={"resource": "train"}, uses=PersistableTestComponent, fn="run_inference", constructor_name="load", config={}, is_target=True, ), }) runner = DaskGraphRunner( graph_schema=graph_schema, model_storage=default_model_storage, execution_context=ExecutionContext(graph_schema=graph_schema, model_id="1"), ) results = runner.run() assert results["load"] == test_value