def get_env(self):
env = super().get_env()
spark_env = {}
if self.get_executor() == "spark":
spark_env = get_mesos_spark_env(
spark_app_name=
"tron_spark_{self.get_service()}_{self.get_instance()}",
spark_ui_port=pick_random_port(
f"{self.get_service()}{self.get_instance()}".encode()),
mesos_leader=find_mesos_leader(self.get_cluster()),
mesos_secret=load_mesos_secret_for_spark(),
paasta_cluster=self.get_cluster(),
paasta_pool=self.get_pool(),
paasta_service=self.get_service(),
paasta_instance=self.get_instance(),
docker_img=self.get_docker_url(),
volumes=format_volumes(
self.get_volumes(
load_system_paasta_config().get_volumes())),
user_spark_opts=self.config_dict.get("spark_args"),
event_log_dir=get_default_event_log_dir(
service=self.get_service(),
aws_credentials_yaml=self.config_dict.get(
"aws_credentials"),
),
)
env["SPARK_OPTS"] = stringify_spark_env(spark_env)
return env
def get_spark_config_dict(self):
if self.get_spark_cluster_manager() == "mesos":
spark_env = get_mesos_spark_env(
spark_app_name=
f"tron_spark_{self.get_service()}_{self.get_instance()}",
spark_ui_port=self.spark_ui_port,
mesos_leader=(find_mesos_leader(
self.get_spark_paasta_cluster())
if not self.for_validation else "N/A"),
paasta_cluster=self.get_spark_paasta_cluster(),
paasta_pool=self.get_spark_paasta_pool(),
paasta_service=self.get_service(),
paasta_instance=self.get_instance(),
docker_img=self.get_docker_url(),
volumes=[
f"{v['hostPath']}:{v['containerPath']}:{v['mode'].lower()}"
for v in self.get_volumes(
load_system_paasta_config().get_volumes())
],
user_spark_opts=self.config_dict.get("spark_args", {}),
event_log_dir=get_default_event_log_dir(
service=self.get_service(),
aws_credentials_yaml=self.config_dict.get(
"aws_credentials_yaml"),
),
needs_docker_cfg=True,
)
else:
spark_env = get_k8s_spark_env(
spark_app_name=
f"tron_spark_{self.get_service()}_{self.get_instance()}",
spark_ui_port=self.spark_ui_port,
paasta_cluster=self.get_spark_paasta_cluster(),
paasta_service=self.get_service(),
paasta_instance=self.get_instance(),
docker_img=self.get_docker_url(),
volumes=self.get_volumes(
load_system_paasta_config().get_volumes()),
user_spark_opts=self.config_dict.get("spark_args", {}),
event_log_dir=get_default_event_log_dir(
service=self.get_service(),
aws_credentials_yaml=self.config_dict.get(
"aws_credentials_yaml"),
),
)
return spark_env
def test_get_default_event_log_dir(self, mock_account_id, account_id, expected_dir):
mock_account_id.return_value = account_id
assert (
get_default_event_log_dir(
access_key="test_access_key", secret_key="test_secret_key"
)
== expected_dir
)
def get_spark_config(
args,
spark_app_name,
spark_ui_port,
docker_img,
system_paasta_config,
volumes,
access_key,
secret_key,
):
# User configurable Spark options
user_args = {
"spark.app.name": spark_app_name,
"spark.cores.max": "4",
"spark.executor.cores": "2",
"spark.executor.memory": "4g",
# Use \; for multiple constraints. e.g.
# instance_type:m4.10xlarge\;pool:default
"spark.mesos.constraints": "pool:%s" % args.pool,
"spark.mesos.executor.docker.forcePullImage": "true",
"spark.mesos.role": "spark",
}
default_event_log_dir = get_default_event_log_dir(
access_key=access_key, secret_key=secret_key
)
if default_event_log_dir is not None:
user_args["spark.eventLog.enabled"] = "true"
user_args["spark.eventLog.dir"] = default_event_log_dir
try:
mesos_address = find_mesos_leader(args.cluster)
except MesosLeaderUnavailable as e:
print(
f"Couldn't reach the {args.cluster} Mesos leader from here. Please run this command from the environment that matches {args.cluster}.\nError: {e}",
file=sys.stderr,
)
sys.exit(2)
# Spark options managed by PaaSTA
paasta_instance = get_smart_paasta_instance_name(args)
non_user_args = {
"spark.master": "mesos://%s" % mesos_address,
"spark.ui.port": spark_ui_port,
"spark.executorEnv.PAASTA_SERVICE": args.service,
"spark.executorEnv.PAASTA_INSTANCE": paasta_instance,
"spark.executorEnv.PAASTA_CLUSTER": args.cluster,
"spark.executorEnv.PAASTA_INSTANCE_TYPE": "spark",
"spark.mesos.executor.docker.parameters": f"label=paasta_service={args.service},label=paasta_instance={paasta_instance}",
"spark.mesos.executor.docker.volumes": ",".join(volumes),
"spark.mesos.executor.docker.image": docker_img,
"spark.mesos.principal": "spark",
"spark.mesos.secret": load_mesos_secret_for_spark(),
}
if not args.build and not args.image:
non_user_args["spark.mesos.uris"] = "file:///root/.dockercfg"
if args.spark_args:
spark_args = args.spark_args.split()
for spark_arg in spark_args:
fields = spark_arg.split("=", 1)
if len(fields) != 2:
print(
PaastaColors.red(
"Spark option %s is not in format option=value." % spark_arg
),
file=sys.stderr,
)
sys.exit(1)
if fields[0] in non_user_args:
print(
PaastaColors.red(
"Spark option {} is set by PaaSTA with {}.".format(
fields[0], non_user_args[fields[0]]
)
),
file=sys.stderr,
)
sys.exit(1)
# Update default configuration
user_args[fields[0]] = fields[1]
if "spark.sql.shuffle.partitions" not in user_args:
num_partitions = str(2 * int(user_args["spark.cores.max"]))
user_args["spark.sql.shuffle.partitions"] = num_partitions
print(
PaastaColors.yellow(
f"Warning: spark.sql.shuffle.partitions has been set to"
f" {num_partitions} to be equal to twice the number of "
f"requested cores, but you should consider setting a "
f"higher value if necessary."
)
)
if int(user_args["spark.cores.max"]) < int(user_args["spark.executor.cores"]):
print(
PaastaColors.red(
"Total number of cores {} is less than per-executor cores {}.".format(
user_args["spark.cores.max"], user_args["spark.executor.cores"]
)
),
file=sys.stderr,
)
sys.exit(1)
exec_mem = user_args["spark.executor.memory"]
if exec_mem[-1] != "g" or not exec_mem[:-1].isdigit() or int(exec_mem[:-1]) > 32:
print(
PaastaColors.red(
"Executor memory {} not in format dg (d<=32).".format(
user_args["spark.executor.memory"]
)
),
file=sys.stderr,
)
sys.exit(1)
# Limit a container's cpu usage
non_user_args["spark.mesos.executor.docker.parameters"] += ",cpus={}".format(
user_args["spark.executor.cores"]
)
return dict(non_user_args, **user_args)