Improve experiment spec validation

CHANGELOG.md

@@ -15,6 +15,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Changed default spec version to `v2`.
 
+### Fixed
+
+- Improved experiment spec validation in `Beaker.experiment.create()` to raise more specific error types.
+
 ## [v0.8.0](https://github.com/allenai/beaker-py/releases/tag/v0.8.0) - 2022-04-12
 
 ### Changed

beaker/services/experiment.py

@@ -32,6 +32,9 @@ def create(
         :raises WorkspaceNotSet: If neither ``workspace`` nor
             :data:`Beaker.config.defeault_workspace <beaker.Config.default_workspace>` are set.
         :raises ImageNotFound: If the image specified by the spec doesn't exist.
+        :raises DatasetNotFound: If a source dataset in the spec doesn't exist.
+        :raises SecretNotFound: If a source secret in the spec doesn't exist.
+        :raises ClusterNotFound: If the cluster in the spec doesn't exist.
         :raises HTTPError: Any other HTTP exception that can occur.
 
         """
@@ -42,9 +45,8 @@ def create(
         else:
             spec = ExperimentSpec.from_file(spec)
             json_spec = spec.to_json()
-        self._validate_spec(spec)
-
         workspace: Workspace = self._resolve_workspace(workspace)
+        self._validate_spec(spec, workspace)
         experiment_data = self.request(
             f"workspaces/{workspace.id}/experiments",
             method="POST",
@@ -304,8 +306,26 @@ def _validate_name(self, name: str) -> None:
         if not name.replace("-", "").replace("_", "").isalnum():
             raise ValueError(err_msg)
 
-    def _validate_spec(self, spec: ExperimentSpec) -> None:
+    def _validate_spec(self, spec: ExperimentSpec, workspace: Workspace) -> None:
         for task in spec.tasks:
             # Make sure image exists.
             if task.image.beaker is not None:
                 self.beaker.image.get(task.image.beaker)
+            # Make sure all beaker data sources exist.
+            for data_mount in task.datasets or []:
+                source = data_mount.source
+                if source.beaker is not None:
+                    self.beaker.dataset.get(source.beaker)
+                if source.secret is not None:
+                    self.beaker.secret.get(source.secret, workspace=workspace)
+                if source.result is not None:
+                    if source.result not in {t.name for t in spec.tasks}:
+                        raise ValueError(
+                            f"Data mount result source '{source.result}' not found in spec"
+                        )
+            # Make sure secrets in env variables exist.
+            for env_var in task.env_vars or []:
+                if env_var.secret is not None:
+                    self.beaker.secret.get(env_var.secret, workspace=workspace)
+            # Make sure cluster exists.
+            self.beaker.cluster.get(task.context.cluster)

tests/experiment_test.py

@@ -2,12 +2,17 @@
 
 from beaker import (
     Beaker,
+    ClusterNotFound,
     CurrentJobStatus,
+    DataMount,
     Dataset,
+    DatasetNotFound,
+    DataSource,
     ExperimentSpec,
     ImageNotFound,
     ImageSource,
     ResultSpec,
+    SecretNotFound,
     Task,
     TaskContext,
     TaskSpec,
@@ -79,3 +84,87 @@ def test_create_experiment_image_not_found(
     )
     with pytest.raises(ImageNotFound):
         client.experiment.create(experiment_name, spec)
+
+
+def test_create_experiment_dataset_not_found(
+    client: Beaker,
+    experiment_name: str,
+    beaker_cluster_name: str,
+):
+    spec = ExperimentSpec(
+        tasks=[
+            TaskSpec(
+                name="main",
+                image=ImageSource(docker="hello-world"),
+                context=TaskContext(cluster=beaker_cluster_name),
+                result=ResultSpec(path="/unused"),
+                datasets=[
+                    DataMount(source=DataSource(beaker="does-not-exist"), mount_path="/data")
+                ],
+            ),
+        ],
+    )
+    with pytest.raises(DatasetNotFound):
+        client.experiment.create(experiment_name, spec)
+
+
+def test_create_experiment_secret_not_found(
+    client: Beaker,
+    experiment_name: str,
+    beaker_cluster_name: str,
+):
+    spec = ExperimentSpec(
+        tasks=[
+            TaskSpec(
+                name="main",
+                image=ImageSource(docker="hello-world"),
+                context=TaskContext(cluster=beaker_cluster_name),
+                result=ResultSpec(path="/unused"),
+                datasets=[
+                    DataMount(source=DataSource(secret="does-not-exist"), mount_path="/data")
+                ],
+            ),
+        ],
+    )
+    with pytest.raises(SecretNotFound):
+        client.experiment.create(experiment_name, spec)
+
+
+def test_create_experiment_result_not_found(
+    client: Beaker,
+    experiment_name: str,
+    beaker_cluster_name: str,
+):
+    spec = ExperimentSpec(
+        tasks=[
+            TaskSpec(
+                name="main",
+                image=ImageSource(docker="hello-world"),
+                context=TaskContext(cluster=beaker_cluster_name),
+                result=ResultSpec(path="/unused"),
+                datasets=[
+                    DataMount(source=DataSource(result="does-not-exist"), mount_path="/data")
+                ],
+            ),
+        ],
+    )
+    with pytest.raises(ValueError, match="does-not-exist"):
+        client.experiment.create(experiment_name, spec)
+
+
+def test_create_experiment_cluster_not_found(
+    client: Beaker,
+    experiment_name: str,
+):
+    spec = ExperimentSpec(
+        tasks=[
+            TaskSpec(
+                name="main",
+                image=ImageSource(docker="hello-world"),
+                context=TaskContext(cluster="does-not-exist"),
+                result=ResultSpec(path="/unused"),
+            ),
+        ],
+    )
+    with pytest.raises(ClusterNotFound):
+        client.experiment.create(experiment_name, spec)