python/pyspark/worker_util.py - spark - Git at Google

 #
 # Licensed to the Apache Software Foundation (ASF) under one or more
 # contributor license agreements.  See the NOTICE file distributed with
 # this work for additional information regarding copyright ownership.
 # The ASF licenses this file to You 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.
 #

 """
 Util functions for workers.
 """

 from contextlib import contextmanager
 import importlib
 from inspect import currentframe, getframeinfo
 import os
 import sys
 from typing import Any, Generator, IO, Optional
 import warnings

 if "SPARK_TESTING" in os.environ:
     assert os.environ.get("SPARK_PYTHON_RUNTIME") == "PYTHON_WORKER", (
         "This module can only be imported in python woker"
     )

 # 'resource' is a Unix specific module.
 has_resource_module = True
 try:
     import resource
 except ImportError:
     has_resource_module = False

 from pyspark.accumulators import _accumulatorRegistry
 from pyspark.util import is_remote_only
 from pyspark.errors import PySparkRuntimeError
 from pyspark.util import local_connect_and_auth
 from pyspark.serializers import (
     read_bool,
     read_int,
     read_long,
     write_int,
     FramedSerializer,
     UTF8Deserializer,
     CPickleSerializer,
 )

 pickleSer = CPickleSerializer()
 utf8_deserializer = UTF8Deserializer()


 def add_path(path: str) -> bool:
     # worker can be used, so do not add path multiple times
     if path not in sys.path:
         # overwrite system packages
         sys.path.insert(1, path)
         return True
     return False


 def read_command(serializer: FramedSerializer, file: IO) -> Any:
     if not is_remote_only():
         from pyspark.core.broadcast import Broadcast

     command = serializer._read_with_length(file)
     if not is_remote_only() and isinstance(command, Broadcast):
         command = serializer.loads(command.value)
     return command


 def check_python_version(infile: IO) -> None:
     """
     Check the Python version between the running process and the one used to serialize the command.
     """
     version = utf8_deserializer.loads(infile)
     worker_version = "%d.%d" % sys.version_info[:2]
     if version != worker_version:
         raise PySparkRuntimeError(
             errorClass="PYTHON_VERSION_MISMATCH",
             messageParameters={
                 "worker_version": worker_version,
                 "driver_version": str(version),
             },
         )


 def setup_memory_limits(memory_limit_mb: int) -> None:
     """
     Sets up the memory limits.

     If memory_limit_mb > 0 and `resource` module is available, sets the memory limit.
     Windows does not support resource limiting and actual resource is not limited on MacOS.
     """
     if memory_limit_mb > 0 and has_resource_module:
         total_memory = resource.RLIMIT_AS
         try:
             soft_limit, hard_limit = resource.getrlimit(total_memory)
             msg = "Current mem limits: {0} of max {1}\n".format(soft_limit, hard_limit)
             print(msg, file=sys.stderr)

             # convert to bytes
             new_limit = memory_limit_mb * 1024 * 1024

             if soft_limit == resource.RLIM_INFINITY or new_limit < soft_limit:
                 msg = "Setting mem limits to {0} of max {1}\n".format(new_limit, new_limit)
                 print(msg, file=sys.stderr)
                 resource.setrlimit(total_memory, (new_limit, new_limit))

         except (resource.error, OSError, ValueError) as e:
             # not all systems support resource limits, so warn instead of failing
             current = currentframe()
             lineno = getframeinfo(current).lineno + 1 if current is not None else 0
             if "__file__" in globals():
                 print(
                     warnings.formatwarning(
                         "Failed to set memory limit: {0}".format(e),
                         ResourceWarning,
                         __file__,
                         lineno,
                     ),
                     file=sys.stderr,
                 )


 def setup_spark_files(infile: IO) -> None:
     """
     Set up Spark files, archives, and pyfiles.
     """
     # fetch name of workdir
     spark_files_dir = utf8_deserializer.loads(infile)

     if not is_remote_only():
         from pyspark.core.files import SparkFiles

         SparkFiles._root_directory = spark_files_dir
         SparkFiles._is_running_on_worker = True

     # fetch names of includes (*.zip and *.egg files) and construct PYTHONPATH
     path_changed = add_path(spark_files_dir)  # *.py files that were added will be copied here
     num_python_includes = read_int(infile)
     for _ in range(num_python_includes):
         filename = utf8_deserializer.loads(infile)
         path_changed = add_path(os.path.join(spark_files_dir, filename)) or path_changed

     if path_changed:
         importlib.invalidate_caches()


 def setup_broadcasts(infile: IO) -> None:
     """
     Set up broadcasted variables.
     """
     if not is_remote_only():
         from pyspark.core.broadcast import Broadcast, _broadcastRegistry

     # fetch names and values of broadcast variables
     needs_broadcast_decryption_server = read_bool(infile)
     num_broadcast_variables = read_int(infile)
     if needs_broadcast_decryption_server:
         # read the decrypted data from a server in the jvm
         conn_info = read_int(infile)
         auth_secret = None
         if conn_info == -1:
             conn_info = utf8_deserializer.loads(infile)
         else:
             auth_secret = utf8_deserializer.loads(infile)
         broadcast_sock_file, _ = local_connect_and_auth(conn_info, auth_secret)

     for _ in range(num_broadcast_variables):
         bid = read_long(infile)
         if bid >= 0:
             if needs_broadcast_decryption_server:
                 read_bid = read_long(broadcast_sock_file)
                 assert read_bid == bid
                 _broadcastRegistry[bid] = Broadcast(sock_file=broadcast_sock_file)
             else:
                 path = utf8_deserializer.loads(infile)
                 _broadcastRegistry[bid] = Broadcast(path=path)

         else:
             bid = -bid - 1
             _broadcastRegistry.pop(bid)

     if needs_broadcast_decryption_server:
         broadcast_sock_file.write(b"1")
         broadcast_sock_file.close()


 @contextmanager
 def get_sock_file_to_executor(timeout: Optional[int] = -1) -> Generator[IO, None, None]:
     # Read information about how to connect back to the JVM from the environment.
     conn_info = os.environ.get(
         "PYTHON_WORKER_FACTORY_SOCK_PATH", int(os.environ.get("PYTHON_WORKER_FACTORY_PORT", -1))
     )
     auth_secret = os.environ.get("PYTHON_WORKER_FACTORY_SECRET")
     sock_file, sock = local_connect_and_auth(conn_info, auth_secret)
     if timeout is None or timeout > 0:
         sock.settimeout(timeout)
     # TODO: Remove the following two lines and use `Process.pid()` when we drop JDK 8.
     write_int(os.getpid(), sock_file)
     sock_file.flush()
     try:
         yield sock_file
     finally:
         sock_file.close()


 def send_accumulator_updates(outfile: IO) -> None:
     """
     Send the accumulator updates back to JVM.
     """
     write_int(len(_accumulatorRegistry), outfile)
     for aid, accum in _accumulatorRegistry.items():
         pickleSer._write_with_length((aid, accum._value), outfile)


 class Conf:
     def __init__(self, infile: Optional[IO] = None) -> None:
         self._conf: dict[str, Any] = {}
         if infile is not None:
             self.load(infile)

     def load(self, infile: IO) -> None:
         num_conf = read_int(infile)
         # We do a sanity check here to reduce the possibility to stuck indefinitely
         # due to an invalid messsage. If the numer of configurations is obviously
         # wrong, we just raise an error directly.
         # We hand-pick the configurations to send to the worker so the number should
         # be very small (less than 100).
         if num_conf < 0 or num_conf > 10000:
             raise PySparkRuntimeError(
                 errorClass="PROTOCOL_ERROR",
                 messageParameters={
                     "failure": f"Invalid number of configurations: {num_conf}",
                 },
             )
         for _ in range(num_conf):
             k = utf8_deserializer.loads(infile)
             v = utf8_deserializer.loads(infile)
             self._conf[k] = v

     def get(self, key: str, default: Any = "", *, lower_str: bool = True) -> Any:
         val = self._conf.get(key, default)
         if isinstance(val, str) and lower_str:
             return val.lower()
         return val
	#
	# Licensed to the Apache Software Foundation (ASF) under one or more
	# contributor license agreements. See the NOTICE file distributed with
	# this work for additional information regarding copyright ownership.
	# The ASF licenses this file to You 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.
	#

	"""
	Util functions for workers.
	"""

	from contextlib import contextmanager
	import importlib
	from inspect import currentframe, getframeinfo
	import os
	import sys
	from typing import Any, Generator, IO, Optional
	import warnings

	if "SPARK_TESTING" in os.environ:
	assert os.environ.get("SPARK_PYTHON_RUNTIME") == "PYTHON_WORKER", (
	"This module can only be imported in python woker"
	)

	# 'resource' is a Unix specific module.
	has_resource_module = True
	try:
	import resource
	except ImportError:
	has_resource_module = False

	from pyspark.accumulators import _accumulatorRegistry
	from pyspark.util import is_remote_only
	from pyspark.errors import PySparkRuntimeError
	from pyspark.util import local_connect_and_auth
	from pyspark.serializers import (
	read_bool,
	read_int,
	read_long,
	write_int,
	FramedSerializer,
	UTF8Deserializer,
	CPickleSerializer,
	)

	pickleSer = CPickleSerializer()
	utf8_deserializer = UTF8Deserializer()


	def add_path(path: str) -> bool:
	# worker can be used, so do not add path multiple times
	if path not in sys.path:
	# overwrite system packages
	sys.path.insert(1, path)
	return True
	return False


	def read_command(serializer: FramedSerializer, file: IO) -> Any:
	if not is_remote_only():
	from pyspark.core.broadcast import Broadcast

	command = serializer._read_with_length(file)
	if not is_remote_only() and isinstance(command, Broadcast):
	command = serializer.loads(command.value)
	return command


	def check_python_version(infile: IO) -> None:
	"""
	Check the Python version between the running process and the one used to serialize the command.
	"""
	version = utf8_deserializer.loads(infile)
	worker_version = "%d.%d" % sys.version_info[:2]
	if version != worker_version:
	raise PySparkRuntimeError(
	errorClass="PYTHON_VERSION_MISMATCH",
	messageParameters={
	"worker_version": worker_version,
	"driver_version": str(version),
	},
	)


	def setup_memory_limits(memory_limit_mb: int) -> None:
	"""
	Sets up the memory limits.

	If memory_limit_mb > 0 and `resource` module is available, sets the memory limit.
	Windows does not support resource limiting and actual resource is not limited on MacOS.
	"""
	if memory_limit_mb > 0 and has_resource_module:
	total_memory = resource.RLIMIT_AS
	try:
	soft_limit, hard_limit = resource.getrlimit(total_memory)
	msg = "Current mem limits: {0} of max {1}\n".format(soft_limit, hard_limit)
	print(msg, file=sys.stderr)

	# convert to bytes
	new_limit = memory_limit_mb * 1024 * 1024

	if soft_limit == resource.RLIM_INFINITY or new_limit < soft_limit:
	msg = "Setting mem limits to {0} of max {1}\n".format(new_limit, new_limit)
	print(msg, file=sys.stderr)
	resource.setrlimit(total_memory, (new_limit, new_limit))

	except (resource.error, OSError, ValueError) as e:
	# not all systems support resource limits, so warn instead of failing
	current = currentframe()
	lineno = getframeinfo(current).lineno + 1 if current is not None else 0
	if "__file__" in globals():
	print(
	warnings.formatwarning(
	"Failed to set memory limit: {0}".format(e),
	ResourceWarning,
	__file__,
	lineno,
	),
	file=sys.stderr,
	)


	def setup_spark_files(infile: IO) -> None:
	"""
	Set up Spark files, archives, and pyfiles.
	"""
	# fetch name of workdir
	spark_files_dir = utf8_deserializer.loads(infile)

	if not is_remote_only():
	from pyspark.core.files import SparkFiles

	SparkFiles._root_directory = spark_files_dir
	SparkFiles._is_running_on_worker = True

	# fetch names of includes (.zip and .egg files) and construct PYTHONPATH
	path_changed = add_path(spark_files_dir) # *.py files that were added will be copied here
	num_python_includes = read_int(infile)
	for _ in range(num_python_includes):
	filename = utf8_deserializer.loads(infile)
	path_changed = add_path(os.path.join(spark_files_dir, filename)) or path_changed

	if path_changed:
	importlib.invalidate_caches()


	def setup_broadcasts(infile: IO) -> None:
	"""
	Set up broadcasted variables.
	"""
	if not is_remote_only():
	from pyspark.core.broadcast import Broadcast, _broadcastRegistry

	# fetch names and values of broadcast variables
	needs_broadcast_decryption_server = read_bool(infile)
	num_broadcast_variables = read_int(infile)
	if needs_broadcast_decryption_server:
	# read the decrypted data from a server in the jvm
	conn_info = read_int(infile)
	auth_secret = None
	if conn_info == -1:
	conn_info = utf8_deserializer.loads(infile)
	else:
	auth_secret = utf8_deserializer.loads(infile)
	broadcast_sock_file, _ = local_connect_and_auth(conn_info, auth_secret)

	for _ in range(num_broadcast_variables):
	bid = read_long(infile)
	if bid >= 0:
	if needs_broadcast_decryption_server:
	read_bid = read_long(broadcast_sock_file)
	assert read_bid == bid
	_broadcastRegistry[bid] = Broadcast(sock_file=broadcast_sock_file)
	else:
	path = utf8_deserializer.loads(infile)
	_broadcastRegistry[bid] = Broadcast(path=path)

	else:
	bid = -bid - 1
	_broadcastRegistry.pop(bid)

	if needs_broadcast_decryption_server:
	broadcast_sock_file.write(b"1")
	broadcast_sock_file.close()


	@contextmanager
	def get_sock_file_to_executor(timeout: Optional[int] = -1) -> Generator[IO, None, None]:
	# Read information about how to connect back to the JVM from the environment.
	conn_info = os.environ.get(
	"PYTHON_WORKER_FACTORY_SOCK_PATH", int(os.environ.get("PYTHON_WORKER_FACTORY_PORT", -1))
	)
	auth_secret = os.environ.get("PYTHON_WORKER_FACTORY_SECRET")
	sock_file, sock = local_connect_and_auth(conn_info, auth_secret)
	if timeout is None or timeout > 0:
	sock.settimeout(timeout)
	# TODO: Remove the following two lines and use `Process.pid()` when we drop JDK 8.
	write_int(os.getpid(), sock_file)
	sock_file.flush()
	try:
	yield sock_file
	finally:
	sock_file.close()


	def send_accumulator_updates(outfile: IO) -> None:
	"""
	Send the accumulator updates back to JVM.
	"""
	write_int(len(_accumulatorRegistry), outfile)
	for aid, accum in _accumulatorRegistry.items():
	pickleSer._write_with_length((aid, accum._value), outfile)


	class Conf:
	def __init__(self, infile: Optional[IO] = None) -> None:
	self._conf: dict[str, Any] = {}
	if infile is not None:
	self.load(infile)

	def load(self, infile: IO) -> None:
	num_conf = read_int(infile)
	# We do a sanity check here to reduce the possibility to stuck indefinitely
	# due to an invalid messsage. If the numer of configurations is obviously
	# wrong, we just raise an error directly.
	# We hand-pick the configurations to send to the worker so the number should
	# be very small (less than 100).
	if num_conf < 0 or num_conf > 10000:
	raise PySparkRuntimeError(
	errorClass="PROTOCOL_ERROR",
	messageParameters={
	"failure": f"Invalid number of configurations: {num_conf}",
	},
	)
	for _ in range(num_conf):
	k = utf8_deserializer.loads(infile)
	v = utf8_deserializer.loads(infile)
	self._conf[k] = v

	def get(self, key: str, default: Any = "", *, lower_str: bool = True) -> Any:
	val = self._conf.get(key, default)
	if isinstance(val, str) and lower_str:
	return val.lower()
	return val