Add check for SDK versus container language/version compatibility

----Release Notes----
[]
-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=120976160

Author: silviulica

google/cloud/dataflow/worker/batchworker.py

@@ -53,6 +53,7 @@
 from google.cloud.dataflow.utils import options
 from google.cloud.dataflow.utils import profiler
 from google.cloud.dataflow.utils import retry
+from google.cloud.dataflow.worker import environment
 from google.cloud.dataflow.worker import executor
 from google.cloud.dataflow.worker import logger
 from google.cloud.dataflow.worker import maptask
@@ -301,6 +302,8 @@ def __init__(self, properties, sdk_pipeline_options):
     self.job_id = properties['job_id']
     self.worker_id = properties['worker_id']
     self.service_path = properties['service_path']
+    # TODO(silviuc): Make sure environment_info_path is always specified.
+    self.environment_info_path = properties.get('environment_info_path', None)
     self.pipeline_options = options.PipelineOptions.from_dictionary(
         sdk_pipeline_options)
     self.capabilities = [self.worker_id, 'remote_source', 'custom_source']
@@ -500,15 +503,25 @@ def run(self):
     # happen before they could be reported to the service. If it is not None,
     # worker will use the first work item to report deferred exceptions and
     # fail eventually.
+    # TODO(silviuc): Add the deferred exception mechanism to streaming worker
     deferred_exception_details = None
 
-    logging.info('Loading main session from the staging area...')
-    try:
-      self._load_main_session(self.local_staging_directory)
-    except Exception:  # pylint: disable=broad-except
-      deferred_exception_details = traceback.format_exc()
-      logging.error('Could not load main session: %s',
-                    deferred_exception_details, exc_info=True)
+    if self.environment_info_path is not None:
+      try:
+        environment.check_sdk_compatibility(self.environment_info_path)
+      except Exception:  # pylint: disable=broad-except
+        deferred_exception_details = traceback.format_exc()
+        logging.error('SDK compatibility check failed: %s',
+                      deferred_exception_details, exc_info=True)
+
+    if deferred_exception_details is None:
+      logging.info('Loading main session from the staging area...')
+      try:
+        self._load_main_session(self.local_staging_directory)
+      except Exception:  # pylint: disable=broad-except
+        deferred_exception_details = traceback.format_exc()
+        logging.error('Could not load main session: %s',
+                      deferred_exception_details, exc_info=True)
 
     # Start status HTTP server thread.
     thread = threading.Thread(target=self.status_server)

google/cloud/dataflow/worker/environment.py

@@ -0,0 +1,55 @@
+# Copyright 2016 Google Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Python Dataflow worker environment compatiblity checking."""
+
+import json
+import logging
+
+from google.cloud.dataflow import version
+
+
+def check_sdk_compatibility(environment_info_path):
+  """Checks if the SDK is compatible with the container in which it runs.
+
+  Args:
+    environment_info_path: Path to a file in JSON format. The file is expected
+      to contain a dictionary with at least two properties: 'language'
+      and 'version'.
+
+  Raises:
+    RuntimeError: For version or language mismatches.
+
+  Other exceptions can be raised if the environment file is not present or does
+  not have the right contents. This can happen only if the base container was
+  not built correctly.
+  """
+  logging.info('Checking if container and SDK language and versions match ...')
+  with open(environment_info_path) as f:
+    info = json.loads(f.read())
+  if info['language'] != 'python':
+    message = (
+        'SDK language \'python\' does not match container language \'%s\'. '
+        'Please rebuild the container using a matching language container.' % (
+            info['language']))
+    logging.error(message)
+    raise RuntimeError(message)
+  if info['version'] != version.__version__:
+    message = (
+        'SDK version %s does not match container version %s. '
+        'Please rebuild the container or use a matching version '
+        'of the SDK.' % (
+            version.__version__, info['version']))
+    logging.error(message)
+    raise RuntimeError(message)

google/cloud/dataflow/worker/environment_test.py

@@ -0,0 +1,70 @@
+# Copyright 2016 Google Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for worker environment compatibility checking."""
+
+import logging
+import os
+import tempfile
+import unittest
+
+from google.cloud.dataflow import version
+from google.cloud.dataflow.worker import environment
+
+
+class EnvironmentTest(unittest.TestCase):
+
+  def create_temp_file(self, path, contents):
+    with open(path, 'w') as f:
+      f.write(contents)
+      return f.name
+
+  def test_basics(self):
+    config_path = os.path.join(tempfile.mkdtemp(), 'config')
+    self.create_temp_file(
+        config_path,
+        '{"language":"python", "version": "%s"}' % version.__version__)
+    environment.check_sdk_compatibility(config_path)
+    # If we get here the test passes since no exception was raised.
+
+  def test_language_no_match(self):
+    config_path = os.path.join(tempfile.mkdtemp(), 'config')
+    self.create_temp_file(
+        config_path,
+        '{"language":"java", "version": "%s"}' % version.__version__)
+    with self.assertRaises(RuntimeError) as exn:
+      environment.check_sdk_compatibility(config_path)
+    self.assertEqual(
+        'SDK language \'python\' does not match container language \'java\'. '
+        'Please rebuild the container using a matching language container.',
+        exn.exception.message)
+
+  def test_version_no_match(self):
+    config_path = os.path.join(tempfile.mkdtemp(), 'config')
+    self.create_temp_file(
+        config_path, '{"language":"python", "version": "0.0.0"}')
+    with self.assertRaises(RuntimeError) as exn:
+      environment.check_sdk_compatibility(config_path)
+    self.assertEqual(
+        'SDK version %s does not match container version 0.0.0. '
+        'Please rebuild the container or use a matching version '
+        'of the SDK.' % (
+            version.__version__),
+        exn.exception.message)
+
+
+if __name__ == '__main__':
+  logging.getLogger().setLevel(logging.INFO)
+  unittest.main()
+