Repository: bacalhau-project/bacalhau (language: python)

Issue:
WASM executor fails to load module dependencies if non-WASM input data is specified
https://github.com/bacalhau-project/bacalhau/blob/a68357cb1efa9ba908f4f99bffa184e505378a00/pkg/executor/wasm/loader.go#L177-L180

If a user specifies a job spec which uses the WASM module depdency feature to load in WASM from multiple modules, it will fail if there is more than a single dependent module or if the WASM module to load is not specified first in the input sources list. 


Discussion / hints:

Fix patch:
diff --git a/pkg/executor/wasm/loader.go b/pkg/executor/wasm/loader.go
index 252e5379ca..a82df396ef 100644
--- a/pkg/executor/wasm/loader.go
+++ b/pkg/executor/wasm/loader.go
@@ -7,7 +7,6 @@ import (
 	"path/filepath"
 	"sync"
 
-	"github.com/bacalhau-project/bacalhau/pkg/models"
 	"github.com/rs/zerolog/log"
 	"github.com/tetratelabs/wazero"
 	"github.com/tetratelabs/wazero/api"
@@ -145,9 +144,9 @@ func (loader *ModuleLoader) InstantiateRemoteModule(ctx context.Context, m stora
 	return loader.runtime.InstantiateModule(ctx, module, loader.config.WithName(m.InputSource.Alias))
 }
 
-const unknownModuleErrStr = ("unknown WASM module with name %q. " +
+const unknownModuleErrStr = ("could not find WASM module with name %q. " +
 	"expecting a built-in module name (like %q), " +
-	"or the name of a module specified as a job input source. " +
+	"or the Alias of one of the job's InputSources. " +
 	"see also: https://docs.bacalhau.org/getting-started/wasm-workload-onboarding")
 
 func (loader *ModuleLoader) loadModuleByName(ctx context.Context, moduleName string) (api.Module, error) {
@@ -179,7 +178,7 @@ func (loader *ModuleLoader) loadModuleByName(ctx context.Context, moduleName str
 	// module, or if the first storage is a WASM module but not the one we were
 	// looking for?
 	for _, s := range loader.storages {
-		if s.InputSource.Source.IsType(models.StorageSourceIPFS) || s.InputSource.Source.IsType(models.StorageSourceURL) {
+		if s.InputSource.Alias == moduleName {
 			return loader.InstantiateRemoteModule(ctx, s)
 		}
 	}
diff --git a/pkg/storage/inline/storage.go b/pkg/storage/inline/storage.go
index bac6727a13..927110ae63 100644
--- a/pkg/storage/inline/storage.go
+++ b/pkg/storage/inline/storage.go
@@ -180,4 +180,15 @@ func (*InlineStorage) Upload(ctx context.Context, path string) (models.SpecConfi
 	}, nil
 }
 
+// StoreBytes returns the passed data embedded as a "data:" URL in a SpecConfig.
+// The input is never compressed.
+func (*InlineStorage) StoreBytes(data []byte) models.SpecConfig {
+	return models.SpecConfig{
+		Type: models.StorageSourceInline,
+		Params: Source{
+			URL: dataurl.EncodeBytes(data),
+		}.ToMap(),
+	}
+}
+
 var _ storage.Storage = (*InlineStorage)(nil)


Test patch:
diff --git a/pkg/executor/wasm/loader_test.go b/pkg/executor/wasm/loader_test.go
new file mode 100644
index 0000000000..79df3a36ee
--- /dev/null
+++ b/pkg/executor/wasm/loader_test.go
@@ -0,0 +1,97 @@
+//go:build unit || !integration
+
+package wasm
+
+import (
+	"context"
+	"testing"
+
+	"github.com/bacalhau-project/bacalhau/pkg/logger"
+	"github.com/bacalhau-project/bacalhau/pkg/models"
+	"github.com/bacalhau-project/bacalhau/pkg/storage"
+	"github.com/bacalhau-project/bacalhau/pkg/storage/inline"
+	"github.com/bacalhau-project/bacalhau/testdata/wasm/dynamic"
+	"github.com/bacalhau-project/bacalhau/testdata/wasm/easter"
+	"github.com/bacalhau-project/bacalhau/testdata/wasm/noop"
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/require"
+	"github.com/tetratelabs/wazero"
+)
+
+func prepareModule(t *testing.T, alias string, program []byte) storage.PreparedStorage {
+	store := inline.NewStorage()
+	spec := store.StoreBytes(program)
+	inputSource := models.InputSource{Source: &spec, Alias: alias, Target: "/" + uuid.NewString()}
+	preparedVolume, err := store.PrepareStorage(context.Background(), t.TempDir(), inputSource)
+	require.NoError(t, err)
+
+	return storage.PreparedStorage{
+		InputSource: inputSource,
+		Volume:      preparedVolume,
+	}
+}
+
+func TestModuleLoading(t *testing.T) {
+	logger.ConfigureTestLogging(t)
+
+	testCases := []struct {
+		name          string
+		entryModule   storage.PreparedStorage
+		importModules []storage.PreparedStorage
+		errorChecker  func(t require.TestingT, err error, _ ...any)
+		moduleChecker func(t require.TestingT, object any, _ ...any)
+	}{
+		{
+			name:          "simple load of single entry module",
+			entryModule:   prepareModule(t, "", noop.Program()),
+			importModules: []storage.PreparedStorage{},
+			errorChecker:  require.NoError,
+			moduleChecker: require.NotNil,
+		},
+		{
+			name:        "does not attempt to load every input source as a WASM module",
+			entryModule: prepareModule(t, "", noop.Program()),
+			importModules: []storage.PreparedStorage{
+				prepareModule(t, "something", []byte{0x0, 0x1, 0x2}),
+			},
+			errorChecker:  require.NoError,
+			moduleChecker: require.NotNil,
+		},
+		{
+			name:          "fails to load a module if it has a dependency that is not given",
+			entryModule:   prepareModule(t, "", dynamic.Program()),
+			importModules: []storage.PreparedStorage{},
+			errorChecker:  require.Error,
+			moduleChecker: require.Nil,
+		},
+		{
+			name:        "loads module with a dependency if it is specified as an input source with the correct alias",
+			entryModule: prepareModule(t, "", dynamic.Program()),
+			importModules: []storage.PreparedStorage{
+				// The alias being a CID is not relevant for this test – this is
+				// just the "module name" that the `dynamic` program is looking
+				// for in its WASM import header.
+				prepareModule(t, "QmPympgyrEGEdSJ93rqvQkR71QLuQGdhKQtYztFwxpQsid", easter.Program()),
+			},
+			errorChecker:  require.NoError,
+			moduleChecker: require.NotNil,
+		},
+	}
+
+	for _, testCase := range testCases {
+		t.Run(testCase.name, func(t *testing.T) {
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
+			runtime := wazero.NewRuntime(ctx)
+			config := wazero.NewModuleConfig()
+
+			loader := NewModuleLoader(runtime, config, append(testCase.importModules, testCase.entryModule)...)
+			require.NotNil(t, loader)
+
+			module, err := loader.InstantiateRemoteModule(ctx, testCase.entryModule)
+			testCase.errorChecker(t, err)
+			testCase.moduleChecker(t, module)
+		})
+	}
+}
diff --git a/pkg/test/scenario/test_scenarios.go b/pkg/test/scenario/test_scenarios.go
index 952763c382..633355d1d9 100644
--- a/pkg/test/scenario/test_scenarios.go
+++ b/pkg/test/scenario/test_scenarios.go
@@ -229,9 +229,12 @@ func WasmCsvTransform(t testing.TB) Scenario {
 
 func WasmDynamicLink(t testing.TB) Scenario {
 	return Scenario{
-		Inputs: StoredFile(
-			"../../../testdata/wasm/easter/main.wasm",
-			"/inputs",
+		Inputs: ManyStores(
+			StoredText("unused input", "/data"),
+			StoredFile(
+				"../../../testdata/wasm/easter/main.wasm",
+				"/inputs",
+			),
 		),
 		ResultsChecker: FileEquals(
 			downloader.DownloadFilenameStdout,