update build_model and tests

Author: FynnBe

bioimageio/core/build_spec/build_model.py

@@ -12,8 +12,8 @@
 import bioimageio.spec.model as model_spec
 from bioimageio.core import export_resource_package, load_raw_resource_description
 from bioimageio.core.resource_io.nodes import URI
-from bioimageio.core.resource_io.utils import resolve_local_source, resolve_source
 from bioimageio.spec.shared.raw_nodes import ImportableModule, ImportableSourceFile
+from bioimageio.spec.shared.utils import resolve_local_source, resolve_source
 
 try:
     from typing import get_args
@@ -81,6 +81,7 @@ def _get_weights(
     model_kwargs=None,
     tensorflow_version=None,
     opset_version=None,
+    dependencies=None,
     **kwargs,
 ):
     weight_path = resolve_source(original_weight_source, root)
@@ -100,6 +101,8 @@ def _get_weights(
         weights = model_spec.raw_nodes.PytorchStateDictWeightsEntry(
             source=weight_source, sha256=weight_hash, **weight_kwargs
         )
+        if dependencies is not None:
+            weight_kwargs["dependencies"] = _get_dependencies(dependencies, root)
 
     elif weight_type == "onnx":
         if opset_version is None:
@@ -745,6 +748,7 @@ def build_model(
         model_kwargs,
         tensorflow_version=tensorflow_version,
         opset_version=opset_version,
+        dependencies=dependencies,
         **weight_kwargs,
     )
 
@@ -813,8 +817,7 @@ def build_model(
 
     if attachments is not None:
         kwargs["attachments"] = spec.rdf.raw_nodes.Attachments(**attachments)
-    if dependencies is not None:
-        kwargs["dependencies"] = _get_dependencies(dependencies, root)
+
     if maintainers is not None:
         kwargs["maintainers"] = [model_spec.raw_nodes.Maintainer(**m) for m in maintainers]
     if parent is not None:

tests/build_spec/test_build_spec.py

@@ -19,13 +19,14 @@ def _test_build_spec(
 ):
     from bioimageio.core.build_spec import build_model
 
-    model_spec = load_raw_resource_description(spec_path)
+    model_spec = load_raw_resource_description(spec_path, update_to_format="latest")
     root = model_spec.root_path
     assert isinstance(model_spec, spec.model.raw_nodes.Model)
     weight_source = model_spec.weights[weight_type].source
 
     cite = {entry.text: entry.doi if entry.url is missing else entry.url for entry in model_spec.cite}
 
+    dep_file = None
     if weight_type == "pytorch_state_dict":
         weight_spec = model_spec.weights["pytorch_state_dict"]
         model_kwargs = None if weight_spec.kwargs is missing else weight_spec.kwargs
@@ -35,6 +36,7 @@ def _test_build_spec(
             arch_path = os.path.abspath(os.path.join(root, arch_path))
             assert os.path.exists(arch_path)
             architecture = f"{arch_path}:{cls_name}"
+        dep_file = None if weight_spec.dependencies is missing else resolve_source(weight_spec.dependencies.file, root)
         weight_type_ = None  # the weight type can be auto-detected
     elif weight_type == "torchscript":
         architecture = None
@@ -45,7 +47,6 @@ def _test_build_spec(
         model_kwargs = None
         weight_type_ = None  # the weight type can be auto-detected
 
-    dep_file = None if model_spec.dependencies is missing else resolve_source(model_spec.dependencies.file, root)
     authors = [{"name": auth.name, "affiliation": auth.affiliation} for auth in model_spec.authors]
 
     input_axes = [input_.axes for input_ in model_spec.inputs]

tests/resource_io/test_load_rdf.py

@@ -1,10 +1,8 @@
 import os.path
 import pathlib
 from pathlib import Path
-from tempfile import NamedTemporaryFile
 
 import pytest
-from marshmallow import ValidationError
 
 from bioimageio.core.resource_io.utils import resolve_source
 
@@ -80,7 +78,7 @@ def test_load_remote_model_with_folders():
 
     # todo: point to real model with nested folders, not this temporary sandbox one
     rdf_url = "https://sandbox.zenodo.org/record/892199/files/rdf.yaml"
-    raw_model = load_raw_resource_description(rdf_url)
+    raw_model = load_raw_resource_description(rdf_url, update_to_format="latest")
     assert isinstance(raw_model, raw_nodes.Model)
     model = load_resource_description(rdf_url)
     assert isinstance(model, nodes.Model)