import os
import pymongo
from urllib.parse import urlsplit
from mspasspy.db.client import DBClient
from mspasspy.util.db_utils import MongoDBWorker
from mspasspy.db.database import Database
from mspasspy.global_history.manager import GlobalHistoryManager
try:
from pyspark import SparkConf, SparkContext
from pyspark.sql import SparkSession
except Exception as err:
SparkConf = None
SparkContext = None
SparkSession = None
_mspasspy_has_pyspark = False
_mspasspy_pyspark_import_error = err
else:
_mspasspy_has_pyspark = True
_mspasspy_pyspark_import_error = None
try:
from dask.distributed import Client as DaskClient
_mspasspy_has_dask_distributed = True
except ImportError:
_mspasspy_has_dask_distributed = False
from mspasspy.ccore.utility import MsPASSError
def _require_pyspark():
if _mspasspy_has_pyspark:
return
message = "Spark scheduler was requested, but PySpark could not be imported"
if _mspasspy_pyspark_import_error is not None:
message += ": " + str(_mspasspy_pyspark_import_error)
raise MsPASSError(message + ".", "Fatal")
def _address_has_port(address):
if "://" in address:
parsed_address = urlsplit(address)
else:
parsed_address = urlsplit("//" + address)
try:
return parsed_address.port is not None
except ValueError:
return False
def _build_dask_scheduler_address(scheduler_address, scheduler_port=None):
if _address_has_port(scheduler_address):
return scheduler_address
if scheduler_port is None or scheduler_port == "":
scheduler_port = "8786"
return scheduler_address + ":" + str(scheduler_port)
[docs]
class Client:
"""
A client-side representation of MSPASS.
This is the only client users should use in MSPASS. The client manages all the other clients or instances.
It creates and manages a Database client.
It creates and manages a Global Hisotry Manager.
It creates and manages a scheduler(spark/dask)
For the address and port of each client/instances, we first check the user specified parameters, if not then
serach the environment varibales values, if not againm then use the default settings.
An existing :class:`dask.distributed.Client` can be supplied with ``dask_client``.
This is useful for externally managed Dask clusters, including Dask Gateway
clusters. The caller owns the external cluster and should keep it alive while
the MsPASS client is using it.
"""
def __init__(
self,
database_host=None,
scheduler=None,
scheduler_host=None,
job_name="mspass",
database_name="mspass",
schema=None,
collection=None,
dask_client=None,
):
# job_name should be a string
if database_host is not None and not type(database_host) is str:
raise MsPASSError(
"database_host should be a string but "
+ str(type(database_host))
+ " is found.",
"Fatal",
)
if scheduler is not None and scheduler not in ("dask", "spark", "none"):
raise MsPASSError(
"scheduler should be dask, spark, or none but "
+ str(scheduler)
+ " is found.",
"Fatal",
)
if scheduler_host is not None and not type(scheduler_host) is str:
raise MsPASSError(
"scheduler_host should be a string but "
+ str(type(scheduler_host))
+ " is found.",
"Fatal",
)
if job_name is not None and not type(job_name) is str:
raise MsPASSError(
"job_name should be a string but " + str(type(job_name)) + " is found.",
"Fatal",
)
if database_name is not None and not type(database_name) is str:
raise MsPASSError(
"database_name should be a string but "
+ str(type(database_name))
+ " is found.",
"Fatal",
)
# collection should be a string
if collection is not None and type(collection) is not str:
raise MsPASSError(
"collection should be a string but "
+ str(type(collection))
+ " is found.",
"Fatal",
)
if dask_client is not None:
if scheduler == "none":
raise MsPASSError(
"dask_client cannot be used when scheduler is none.",
"Fatal",
)
if scheduler == "spark":
raise MsPASSError(
"dask_client can only be used with the dask scheduler.",
"Fatal",
)
if not _mspasspy_has_dask_distributed or not isinstance(
dask_client, DaskClient
):
raise MsPASSError(
"dask_client should be a dask.distributed.Client but "
+ str(type(dask_client))
+ " is found.",
"Fatal",
)
# check env variables
MSPASS_DB_ADDRESS = os.environ.get("MSPASS_DB_ADDRESS")
MONGODB_PORT = os.environ.get("MONGODB_PORT")
MSPASS_SCHEDULER = os.environ.get("MSPASS_SCHEDULER")
MSPASS_SCHEDULER_ADDRESS = os.environ.get("MSPASS_SCHEDULER_ADDRESS")
DASK_SCHEDULER_PORT = os.environ.get("DASK_SCHEDULER_PORT")
SPARK_MASTER_PORT = os.environ.get("SPARK_MASTER_PORT")
if (
dask_client is None
and not _mspasspy_has_pyspark
and (
scheduler == "spark"
or (scheduler is None and MSPASS_SCHEDULER == "spark")
)
):
_require_pyspark()
# create a database client
# priority: parameter -> env -> default
database_host_has_port = False
if database_host:
database_address = database_host
# check if database_host contains port number already
if ":" in database_address:
database_host_has_port = True
elif MSPASS_DB_ADDRESS:
database_address = MSPASS_DB_ADDRESS
else:
database_address = "127.0.0.1"
# add port
if not database_host_has_port and MONGODB_PORT:
database_address += ":" + MONGODB_PORT
try:
self._db_client = DBClient(database_address)
self._db_client.server_info()
except Exception as err:
raise MsPASSError(
"Runntime error: cannot create a database client with: "
+ database_address,
"Fatal",
)
# set default database name
self._default_database_name = database_name
self._default_schema = schema
self._default_collection = collection
# create a Global History Manager
if schema:
global_history_manager_db = Database(
self._db_client, database_name, db_schema=schema
)
else:
global_history_manager_db = Database(self._db_client, database_name)
self._global_history_manager = GlobalHistoryManager(
global_history_manager_db, job_name, collection=collection
)
# set scheduler
self._scheduler_disabled = False
if dask_client is not None:
self._scheduler = "dask"
elif scheduler:
if scheduler == "none":
self._scheduler = None
self._scheduler_disabled = True
else:
self._scheduler = scheduler
elif MSPASS_SCHEDULER:
if MSPASS_SCHEDULER not in ("dask", "spark", "none"):
raise MsPASSError(
"MSPASS_SCHEDULER should be dask, spark, or none but "
+ str(MSPASS_SCHEDULER)
+ " is found.",
"Fatal",
)
if MSPASS_SCHEDULER == "none":
self._scheduler = None
self._scheduler_disabled = True
else:
self._scheduler = MSPASS_SCHEDULER
else:
if _mspasspy_has_dask_distributed:
self._scheduler = "dask"
elif _mspasspy_has_pyspark:
self._scheduler = "spark"
else:
self._scheduler = None
# scheduler configuration
if self._scheduler == "spark":
scheduler_host_has_port = False
if scheduler_host:
self._spark_master_url = scheduler_host
# add spark:// prefix if not exist
if "spark://" not in scheduler_host:
self._spark_master_url = "spark://" + self._spark_master_url
# check if spark host address contains port number already
if self._spark_master_url.count(":") == 2:
scheduler_host_has_port = True
elif MSPASS_SCHEDULER_ADDRESS:
self._spark_master_url = MSPASS_SCHEDULER_ADDRESS
# add spark:// prefix if not exist
if "spark://" not in MSPASS_SCHEDULER_ADDRESS:
self._spark_master_url = "spark://" + self._spark_master_url
else:
self._spark_master_url = "local"
# add port number
# 1. not the default 'local'
# 2. scheduler_host and does not contain port number
# 3. SPARK_MASTER_PORT exists
if (
(scheduler_host or MSPASS_SCHEDULER_ADDRESS)
and not scheduler_host_has_port
and SPARK_MASTER_PORT
):
self._spark_master_url += ":" + SPARK_MASTER_PORT
# sanity check
try:
spark = (
SparkSession.builder.appName("mspass")
.master(self._spark_master_url)
.getOrCreate()
)
self._spark_context = spark.sparkContext
except Exception as err:
raise MsPASSError(
"Runntime error: cannot create a spark configuration with: "
+ self._spark_master_url,
"Fatal",
)
elif self._scheduler == "dask":
# if no defind scheduler_host and no MSPASS_SCHEDULER_ADDRESS, use local cluster to create a client
if dask_client is not None:
self._dask_client = dask_client
elif not scheduler_host and not MSPASS_SCHEDULER_ADDRESS:
self._dask_client = DaskClient()
else:
if scheduler_host:
scheduler_address = scheduler_host
else:
scheduler_address = MSPASS_SCHEDULER_ADDRESS
self._dask_client_address = _build_dask_scheduler_address(
scheduler_address, DASK_SCHEDULER_PORT
)
# sanity check
try:
self._dask_client = DaskClient(self._dask_client_address)
except Exception as err:
raise MsPASSError(
"Runntime error: cannot create a dask client with: "
+ self._dask_client_address,
"Fatal",
)
elif not self._scheduler_disabled:
print("There is no spark or dask installed, this client has no scheduler")
# Auto-register MongoDB worker plugin for dask to avoid DB serialization leaks
if self._scheduler == "dask":
mongo_plugin = MongoDBWorker(self, dbclient_key="dbclient")
self._dask_client.register_plugin(mongo_plugin, name="mongodb_worker")
[docs]
def get_database_client(self):
"""
Get the database client in the global history manager
:return: :class:`mspasspy.db.database.Database`
"""
return self._db_client
[docs]
def get_database(self, database_name=None):
"""
Get a database by database_name, if database_name is not specified, use the default one
:param database_name: the name of database
:type database_name: :class:`str`
:return: :class:`mspasspy.db.database.Database`
"""
if not database_name:
return Database(self._db_client, self._default_database_name)
return Database(self._db_client, database_name)
[docs]
def get_global_history_manager(self):
"""
Get the global history manager with this client
:return: :class:`mspasspy.global_history.manager.GlobalHistoryManager`
"""
return self._global_history_manager
[docs]
def get_scheduler(self):
"""
Get the scheduler(spark/dask) with this client
:return: :class:`pyspark.SparkContext`/:class:`dask.distributed.Client`/None
"""
if self._scheduler == "spark":
return self._spark_context
elif self._scheduler == "dask":
return self._dask_client
elif self._scheduler_disabled:
return None
else:
print(
"There is no spark or dask installed, this client has no scheduler, returned None"
)
return None
[docs]
def set_database_client(self, database_host, database_port=None):
"""
Set a database client by database_host(and database_port)
:param database_host: the host address of database client
:type database_host: :class:`str`
:param database_port: the port of database client
:type database_port: :class:`str`
"""
database_host_has_port = False
database_address = database_host
# check if port is already in the database_host address
if ":" in database_host:
database_host_has_port = True
# add port
if not database_host_has_port and database_port:
database_address += ":" + database_port
# sanity check
temp_db_client = self._db_client
try:
self._db_client = DBClient(database_address)
self._db_client.server_info()
except Exception as err:
# restore the _db_client
self._db_client = temp_db_client
raise MsPASSError(
"Runntime error: cannot create a database client with: "
+ database_address,
"Fatal",
)
[docs]
def set_global_history_manager(self, history_db, job_name, collection=None):
"""
Set a global history manager by history_db, job_name(and collection)
:param history_db: the database will be set in the global history manager
:type history_db: :class:`mspasspy.db.database.Database`
:param job_name: the job name will be set in the global history manager
:type job_name: :class:`str`
:param collection: the collection name will be set in the history_db
:type collection: :class:`str`
"""
if not isinstance(history_db, Database):
raise TypeError(
"history_db should be a mspasspy.db.Database but "
+ str(type(history_db))
+ " is found."
)
if not type(job_name) is str:
raise TypeError(
"job_name should be a string but " + str(type(job_name)) + " is found."
)
if collection is not None and type(collection) is not str:
raise TypeError(
"collection should be a string but "
+ str(type(collection))
+ " is found."
)
self._global_history_manager = GlobalHistoryManager(
history_db, job_name, collection=collection
)
[docs]
def set_scheduler(self, scheduler, scheduler_host, scheduler_port=None):
"""
Set a scheduler by scheduler type, scheduler_host(and scheduler_port)
:param scheduler: the scheduler type, should be either dask or spark
:type scheduler: :class:`str`
:param scheduler_host: the host address of scheduler
:type scheduler_host: :class:`str`
:param scheduler_port: the port of scheduler
:type scheduler_port: :class:`str`
"""
if scheduler != "dask" and scheduler != "spark":
raise MsPASSError(
"scheduler should be either dask or spark but "
+ str(scheduler)
+ " is found.",
"Fatal",
)
if scheduler == "spark":
_require_pyspark()
prev_scheduler = self._scheduler
self._scheduler = scheduler
if scheduler == "spark":
scheduler_host_has_port = False
self._spark_master_url = scheduler_host
# add spark:// prefix if not exist
if "spark://" not in scheduler_host:
self._spark_master_url = "spark://" + self._spark_master_url
# check if spark host address contains port number already
if self._spark_master_url.count(":") == 2:
scheduler_host_has_port = True
# add port
if not scheduler_host_has_port and scheduler_port:
self._spark_master_url += ":" + scheduler_port
# sanity check
prev_spark_context = None
prev_spark_conf = None
if hasattr(self, "_spark_context"):
prev_spark_context = self._spark_context
prev_spark_conf = self._spark_context.getConf()
try:
if hasattr(self, "_spark_context") and isinstance(
self._spark_context, SparkContext
):
# update the confinguration
spark_conf = self._spark_context._conf.setMaster(
self._spark_master_url
)
else:
spark_conf = (
SparkConf()
.setAppName("mspass")
.setMaster(self._spark_master_url)
)
# stop the previous spark context
# FIXME if the new context does not start, we shouldn't stop the previous here.
# if prev_spark_context:
# prev_spark_context.stop()
# create a new spark context -> might cause error so that execute exception code
spark = SparkSession.builder.config(conf=spark_conf).getOrCreate()
self._spark_context = spark.sparkContext
except Exception as err:
# restore the spark context by the previous spark configuration
if prev_spark_conf:
self._spark_context = SparkContext.getOrCreate(conf=prev_spark_conf)
# restore the scheduler type
if self._scheduler == "spark" and prev_scheduler == "dask":
self._scheduler = prev_scheduler
raise MsPASSError(
"Runntime error: cannot create a spark configuration with: "
+ self._spark_master_url,
"Fatal",
)
# close previous dask client if success
if hasattr(self, "_dask_client"):
del self._dask_client
elif scheduler == "dask":
self._dask_client_address = _build_dask_scheduler_address(
scheduler_host, scheduler_port
)
# sanity check
prev_dask_client = None
if hasattr(self, "_dask_client"):
prev_dask_client = self._dask_client
try:
# create a new dask client
self._dask_client = DaskClient(self._dask_client_address)
except Exception as err:
# restore the dask client if exists
if prev_dask_client:
self._dask_client = prev_dask_client
# restore the scheduler type
if self._scheduler == "dask" and prev_scheduler == "spark":
self._scheduler = prev_scheduler
raise MsPASSError(
"Runntime error: cannot create a dask client with: "
+ self._dask_client_address,
"Fatal",
)
# remove previous spark context if success setting new dask client
if hasattr(self, "_spark_context"):
del self._spark_context
