Merge pull request #1010 from PrefectHQ/docker-version

Docker prefect version

CHANGELOG.md

@@ -10,7 +10,7 @@ These changes are available in the [master branch](https://github.com/PrefectHQ/
 
 ### Enhancements
 
-- None
+- Add `prefect_version` kwarg to `Docker` storage for controlling the version of prefect installed into your containers - [#1010](https://github.com/PrefectHQ/prefect/pull/1010)
 
 ### Task Library
 

requirements.txt

@@ -2,7 +2,7 @@ click >= 7.0, < 8.0
 cloudpickle >= 0.6.0, < 0.7
 croniter >= 0.3, < 0.4
 cryptography >= 2.2.2, < 3.0
-dask >= 0.18, < 2.0
+dask >= 0.18, < 1.2.2
 distributed >= 1.26.1, < 2.0
 docker >= 3.4.1, < 4.0
 idna < 2.8, >= 2.5

src/prefect/environments/storage/docker.py

@@ -36,6 +36,8 @@ class Docker(Storage):
         - env_vars (dict, optional): a dictionary of environment variables to use when building
         - files (dict, optional): a dictionary of files to copy into the image when building
         - base_url: (str, optional): a URL of a Docker daemon to use when for Docker related functionality
+        - prefect_version (str, optional): an optional branch, tag, or commit specifying the version of prefect
+            you want installed into the container; defaults to `"0.5.3"`
     """
 
     def __init__(
@@ -48,6 +50,7 @@ def __init__(
         env_vars: dict = None,
         files: dict = None,
         base_url: str = "unix://var/run/docker.sock",
+        prefect_version: str = None,
     ) -> None:
         self.registry_url = registry_url
         self.base_image = base_image
@@ -59,6 +62,7 @@ def __init__(
         self.flows = dict()  # type: Dict[str, str]
         self._flows = dict()  # type: Dict[str, "prefect.core.flow.Flow"]
         self.base_url = base_url
+        self.prefect_version = prefect_version or "0.5.3"
 
         not_absolute = [
             file_path for file_path in self.files if not os.path.isabs(file_path)
@@ -334,7 +338,7 @@ def create_dockerfile_object(self, directory: str = None) -> None:
                 ENV PREFECT__USER_CONFIG_PATH="/root/.prefect/config.toml"
                 {env_vars}
 
-                RUN pip install git+https://github.com/PrefectHQ/prefect.git@master#egg=prefect[kubernetes]
+                RUN pip install git+https://github.com/PrefectHQ/prefect.git@{version}#egg=prefect[kubernetes]
                 # RUN pip install prefect
 
                 RUN python /root/.prefect/healthcheck.py
@@ -344,6 +348,7 @@ def create_dockerfile_object(self, directory: str = None) -> None:
                     copy_flows=copy_flows,
                     copy_files=copy_files,
                     env_vars=env_vars,
+                    version=self.prefect_version,
                 )
             )
 

src/prefect/serialization/storage.py

@@ -52,6 +52,7 @@ class Meta:
     image_name = fields.String(allow_none=True)
     image_tag = fields.String(allow_none=True)
     flows = fields.Dict(key=fields.Str(), values=fields.Str())
+    prefect_version = fields.String(allow_none=False)
 
     @post_load
     def create_object(self, data: dict) -> Docker:

tests/environments/storage/test_docker_storage.py

@@ -41,6 +41,7 @@ def test_empty_docker_storage():
     assert not storage.python_dependencies
     assert not storage.env_vars
     assert not storage.files
+    assert storage.prefect_version == "0.5.3"
     assert storage.base_url == "unix://var/run/docker.sock"
 
 
@@ -53,6 +54,7 @@ def test_empty_docker_storage():
         image_tag="test5",
         env_vars={"test": "1"},
         base_url="test_url",
+        prefect_version="master",
     )
 
     assert storage.registry_url == "test1"
@@ -62,6 +64,7 @@ def test_empty_docker_storage():
     assert storage.python_dependencies == ["test"]
     assert storage.env_vars == {"test": "1"}
     assert storage.base_url == "test_url"
+    assert storage.prefect_version == "master"
 
 
 def test_files_not_absolute_path():
@@ -185,6 +188,18 @@ def test_create_dockerfile_from_base_image():
         assert "FROM python:3.6" in output
 
 
+def test_create_dockerfile_from_prefect_version():
+    storage = Docker(prefect_version="master")
+
+    with tempfile.TemporaryDirectory() as tempdir:
+        storage.create_dockerfile_object(directory=tempdir)
+
+        with open(os.path.join(tempdir, "Dockerfile"), "r") as dockerfile:
+            output = dockerfile.read()
+
+        assert "prefect.git@master" in output
+
+
 def test_create_dockerfile_with_weird_flow_name():
     with tempfile.TemporaryDirectory() as tempdir_outside:
 

tests/serialization/test_storage.py

@@ -25,6 +25,7 @@ def test_docker_empty_serialize():
 
     assert serialized
     assert serialized["__version__"] == prefect.__version__
+    assert "prefect_version" in serialized
     assert not serialized["image_name"]
     assert not serialized["image_tag"]
     assert not serialized["registry_url"]
@@ -48,7 +49,9 @@ def test_memory_roundtrip():
 
 
 def test_docker_full_serialize():
-    docker = storage.Docker(registry_url="url", image_name="name", image_tag="tag")
+    docker = storage.Docker(
+        registry_url="url", image_name="name", image_tag="tag", prefect_version="0.5.2"
+    )
     serialized = DockerSchema().dump(docker)
 
     assert serialized
@@ -57,6 +60,7 @@ def test_docker_full_serialize():
     assert serialized["image_tag"] == "tag"
     assert serialized["registry_url"] == "url"
     assert serialized["flows"] == dict()
+    assert serialized["prefect_version"] == "0.5.2"
 
 
 def test_docker_serialize_with_flows():