Allow no beaker or docker image in exp spec (#180)

* Allow no beaker or docker image in exp spec

* missing import

beaker/data_model/experiment_spec.py

@@ -4,6 +4,7 @@
 from pydantic import Field, root_validator, validator
 
 from ..aliases import PathOrStr
+from ..exceptions import *
 from .base import BaseModel
 
 __all__ = [
@@ -44,14 +45,6 @@ class ImageSource(BaseModel, frozen=False):
         be pre-configured to enable access.
     """
 
-    @root_validator(pre=True)
-    def _check_exactly_one_field_set(cls, values: Dict[str, Any]) -> Dict[str, Any]:
-        if (values.get("beaker") is None) == (values.get("docker") is None):
-            raise ValueError(
-                "Exactly one of 'beaker' or 'docker' must be specified for image source"
-            )
-        return values
-
 
 class EnvVar(BaseModel, frozen=False):
     """
@@ -718,3 +711,10 @@ def with_description(self, description: str) -> "ExperimentSpec":
         'Hello, Mars!'
         """
         return self.copy(deep=True, update={"description": description})
+
+    def validate(self):
+        for task in self.tasks:
+            if (task.image.beaker is None) == (task.image.docker is None):
+                raise ExperimentSpecError(
+                    "Exactly one of 'beaker' or 'docker' must be specified for image source"
+                )

beaker/exceptions.py

@@ -230,3 +230,7 @@ class JobFailedError(BeakerError):
 
 class JobTimeoutError(BeakerError, TimeoutError):
     pass
+
+
+class ExperimentSpecError(BeakerError):
+    pass

beaker/services/experiment.py

@@ -83,12 +83,10 @@ def create(
 
         """
         self.validate_beaker_name(name)
-        json_spec: Dict[str, Any]
-        if isinstance(spec, ExperimentSpec):
-            json_spec = spec.to_json()
-        else:
+        if not isinstance(spec, ExperimentSpec):
             spec = ExperimentSpec.from_file(spec)
-            json_spec = spec.to_json()
+        spec.validate()
+        json_spec = spec.to_json()
         workspace = self.resolve_workspace(workspace)
         self._validate_spec(spec, workspace)
         experiment_data = self.request(

tests/data_model_test.py

@@ -6,16 +6,6 @@
 from beaker.exceptions import ValidationError
 
 
-def test_image_source_validation():
-    with pytest.raises(ValidationError, match="'beaker' or 'docker'"):
-        ImageSource()
-
-    with pytest.raises(ValidationError, match="'beaker' or 'docker'"):
-        ImageSource(beaker="foo", docker="bar")
-
-    ImageSource(beaker="foo")
-
-
 def test_data_source_validation():
     with pytest.raises(ValidationError, match="Exactly one"):
         DataSource()