Merge pull request spotify#110 from spotify/fallon/use-dockerhub-images

[cli] Use apache beam images in job create templates

cli/src/klio_cli/commands/job/create.py

@@ -37,9 +37,7 @@
     r"topics/(?P<topic>[a-zA-Z0-9-_.~+%]{3,255})"
 )
 VALID_BEAM_PY_VERSIONS = ["3.5", "3.6", "3.7"]
-VALID_BEAM_PY_VERSIONS_SHORT = [
-    "".join(p.split(".")) for p in VALID_BEAM_PY_VERSIONS
-]
+
 DEFAULTS = {
     "region": "europe-west1",
     "experiments": "beam_fn_api",
@@ -195,17 +193,12 @@ def _parse_python_version(self, python_version):
         if len(python_version) < 2 or len(python_version) > 5:
             raise click.BadParameter(invalid_err_msg)
 
-        # 3.x -> 3x; 3.x.y -> 3x
-        if "." in python_version:
-            python_version = "".join(python_version.split(".")[:2])
+        # keep only the major and minor version information
+        python_version = ".".join(python_version.split(".")[:2])
 
-        if python_version not in VALID_BEAM_PY_VERSIONS_SHORT:
+        if python_version not in VALID_BEAM_PY_VERSIONS:
             raise click.BadParameter(invalid_err_msg)
 
-        # Dataflow's 3.5 image is just "python3" while 3.6 and 3.7 are
-        # "python36" and "python37"
-        if python_version == "35":
-            python_version = "3"
         return python_version
 
     def _get_default_streaming_job_context(self, kwargs):

cli/src/klio_cli/commands/job/utils/templates/dockerfile.tpl

@@ -1,5 +1,5 @@
 ## -*- docker-image-name: "{{ klio.pipeline_options.worker_harness_container_image }}" -*-
-FROM dataflow.gcr.io/v1beta3/python{{ klio.python_version }}-fnapi:2.24.0
+FROM apache/beam_python{{ klio.python_version }}_sdk:2.24.0
 
 WORKDIR /usr/src/app
 {%- if klio.use_fnapi %}

cli/tests/commands/job/test_create.py

@@ -124,7 +124,7 @@ def context():
     return {
         "job_name": "test-job",
         "job_type": "streaming",
-        "python_version": "36",
+        "python_version": "3.6",
         "pipeline_options": {
             "project": "test-gcp-project",
             "region": "europe-west1",
@@ -184,7 +184,7 @@ def default_context():
     return {
         "job_name": "test-job",
         "job_type": "streaming",
-        "python_version": "36",
+        "python_version": "3.6",
         "use_fnapi": False,
         "create_resources": False,
         "pipeline_options": {
@@ -357,7 +357,7 @@ def test_create_dockerfile(use_fnapi, tmpdir, job):
             "worker_harness_container_image": "gcr.io/foo/bar",
             "project": "test-gcp-project",
         },
-        "python_version": "36",
+        "python_version": "3.6",
         "use_fnapi": use_fnapi,
         "create_resources": False,
     }
@@ -419,15 +419,12 @@ def test_validate_region_raises(job):
 @pytest.mark.parametrize(
     "input_version,exp_output_version",
     (
-        ("3.5", "3"),
-        ("3.5.1", "3"),
-        ("35", "3"),
-        ("3.6", "36"),
-        ("3.6.1", "36"),
-        ("36", "36"),
-        ("3.7", "37"),
-        ("3.7.1", "37"),
-        ("37", "37"),
+        ("3.5", "3.5"),
+        ("3.5.1", "3.5"),
+        ("3.6", "3.6"),
+        ("3.6.1", "3.6"),
+        ("3.7", "3.7"),
+        ("3.7.1", "3.7"),
     ),
 )
 def test_parse_python_version(input_version, exp_output_version, job):
@@ -487,7 +484,7 @@ def context_overrides():
         "dependencies": [
             {"job_name": "parent-job", "gcp_project": "parent-gcp-project"}
         ],
-        "python_version": "37",
+        "python_version": "3.7",
         "use_fnapi": "n",
         "create_resources": "n",
     }
@@ -532,7 +529,7 @@ def expected_overrides():
                 {"job_name": "parent-job", "gcp_project": "parent-gcp-project"}
             ],
         },
-        "python_version": "37",
+        "python_version": "3.7",
         "use_fnapi": False,
         "create_resources": False,
         "job_type": "streaming",
@@ -678,7 +675,7 @@ def test_get_context_from_user_inputs(
         32,
         "n1-standard-2",
         "",
-        "36",
+        "3.6",
         "gs://test-gcp-project-dataflow-tmp/test-job/staging",
         "gs://test-gcp-project-dataflow-tmp/test-job/temp",
         "projects/test-parent-gcp-project/topics/test-parent-job-output",
@@ -749,7 +746,7 @@ def test_get_context_from_user_inputs_no_prompts(
     )
 
     expected_overrides["pipeline_options"].pop("project")
-    expected_overrides["python_version"] = "36"
+    expected_overrides["python_version"] = "3.6"
     assert not mock_prompt.call_count
     assert not mock_confirm.call_count
     mock_validate_region.assert_called_once_with("us-central1")
@@ -806,7 +803,7 @@ def test_get_context_from_user_inputs_dependency_settings(
     )
 
     expected_overrides["pipeline_options"].pop("project")
-    expected_overrides["python_version"] = "36"
+    expected_overrides["python_version"] = "3.6"
 
     assert not mock_prompt.call_count
     assert 1 == mock_confirm.call_count

cli/tests/commands/job/utils/fixtures/expected/fnapi/Dockerfile

@@ -1,5 +1,5 @@
 ## -*- docker-image-name: "gcr.io/foo/bar" -*-
-FROM dataflow.gcr.io/v1beta3/python36-fnapi:2.24.0
+FROM apache/beam_python3.6_sdk:2.24.0
 
 WORKDIR /usr/src/app
 RUN mkdir -p /usr/src/config

cli/tests/commands/job/utils/fixtures/expected/no_fnapi/Dockerfile

@@ -1,5 +1,5 @@
 ## -*- docker-image-name: "gcr.io/foo/bar" -*-
-FROM dataflow.gcr.io/v1beta3/python36-fnapi:2.24.0
+FROM apache/beam_python3.6_sdk:2.24.0
 
 WORKDIR /usr/src/app