# ---------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# ---------------------------------------------------------
from typing import Dict, List, Optional, Union
from marshmallow import Schema
from ..._schema import PathAwareSchema
from ...constants._job.job import RestSparkConfKey
from ...entities import Environment, SparkJobEntry
from ...entities._job.parameterized_spark import DUMMY_IMAGE, ParameterizedSpark
from ...entities._job.spark_job_entry_mixin import SparkJobEntryMixin
from .._schema.component import InternalSparkComponentSchema
from ..entities import InternalComponent
from .environment import InternalEnvironment
class InternalSparkComponent(
InternalComponent, ParameterizedSpark, SparkJobEntryMixin
): # pylint: disable=too-many-instance-attributes, too-many-ancestors
"""Internal Spark Component
This class is used to handle internal spark component.
It can be loaded from internal spark component yaml or from rest object of an internal spark component.
But after loaded, its structure will be the same as spark component.
"""
def __init__(
self,
entry: Union[Dict[str, str], SparkJobEntry, None] = None,
py_files: Optional[List[str]] = None,
jars: Optional[List[str]] = None,
files: Optional[List[str]] = None,
archives: Optional[List[str]] = None,
driver_cores: Optional[int] = None,
driver_memory: Optional[str] = None,
executor_cores: Optional[int] = None,
executor_memory: Optional[str] = None,
executor_instances: Optional[int] = None,
dynamic_allocation_enabled: Optional[bool] = None,
dynamic_allocation_min_executors: Optional[int] = None,
dynamic_allocation_max_executors: Optional[int] = None,
conf: Optional[Dict[str, str]] = None,
args: Optional[str] = None,
**kwargs,
):
SparkJobEntryMixin.__init__(self, entry=entry, **kwargs)
# environment.setter has been overridden in ParameterizedSpark, so we need to pop it out here
environment = kwargs.pop("environment", None)
InternalComponent.__init__(self, **kwargs)
# Pop it to avoid passing multiple values for code in ParameterizedSpark.__init__
code = kwargs.pop("code", None)
ParameterizedSpark.__init__(
self,
code=self.base_path,
entry=entry,
py_files=py_files,
jars=jars,
files=files,
archives=archives,
conf=conf,
environment=environment,
args=args,
**kwargs,
)
self.code = code
# For pipeline spark job, we also allow user to set driver_cores, driver_memory and so on by setting conf.
# If root level fields are not set by user, we promote conf setting to root level to facilitate subsequent
# verification. This usually happens when we use to_component(SparkJob) or builder function spark() as a node
# in pipeline sdk
conf = conf or {}
self.driver_cores = driver_cores or conf.get(RestSparkConfKey.DRIVER_CORES, None)
self.driver_memory = driver_memory or conf.get(RestSparkConfKey.DRIVER_MEMORY, None)
self.executor_cores = executor_cores or conf.get(RestSparkConfKey.EXECUTOR_CORES, None)
self.executor_memory = executor_memory or conf.get(RestSparkConfKey.EXECUTOR_MEMORY, None)
self.executor_instances = executor_instances or conf.get(RestSparkConfKey.EXECUTOR_INSTANCES, None)
self.dynamic_allocation_enabled = dynamic_allocation_enabled or conf.get(
RestSparkConfKey.DYNAMIC_ALLOCATION_ENABLED, None
)
self.dynamic_allocation_min_executors = dynamic_allocation_min_executors or conf.get(
RestSparkConfKey.DYNAMIC_ALLOCATION_MIN_EXECUTORS, None
)
self.dynamic_allocation_max_executors = dynamic_allocation_max_executors or conf.get(
RestSparkConfKey.DYNAMIC_ALLOCATION_MAX_EXECUTORS, None
)
self.conf = conf
self.args = args
@classmethod
def _create_schema_for_validation(cls, context) -> Union[PathAwareSchema, Schema]:
return InternalSparkComponentSchema(context=context)
@property # type: ignore[override]
def environment(self) -> Optional[Union[Environment, str]]:
"""Get the environment of the component.
:return: The environment of the component.
:rtype: Optional[Union[Environment, str]]]
"""
if isinstance(self._environment, Environment) and self._environment.image is None:
return Environment(conda_file=self._environment.conda_file, image=DUMMY_IMAGE)
return self._environment
@environment.setter
def environment(self, value):
"""Set the environment of the component.
:param value: The environment of the component.
:type value: Union[str, Environment, dict]
:return: No return
:rtype: None
"""
if value is None or isinstance(value, (str, Environment)):
self._environment = value
elif isinstance(value, dict):
internal_environment = InternalEnvironment(**value)
internal_environment.resolve(self.base_path)
self._environment = Environment(
name=internal_environment.name,
version=internal_environment.version,
)
if internal_environment.conda:
self._environment.conda_file = {
"dependencies": internal_environment.conda[InternalEnvironment.CONDA_DEPENDENCIES]
}
if internal_environment.docker:
self._environment.image = internal_environment.docker["image"]
# we suppose that loaded internal spark component won't be used to create another internal spark component
# so the environment construction here can be simplified
else:
raise ValueError(f"Unsupported environment type: {type(value)}")
@property
def jars(self) -> Optional[List[str]]:
"""Get the jars of the component.
:return: The jars of the component.
:rtype: Optional[List[str]]
"""
return self._jars
@jars.setter
def jars(self, value: Union[str, List[str]]):
"""Set the jars of the component.
:param value: The jars of the component.
:type value: Union[str, List[str]]
:return: No return
:rtype: None
"""
if isinstance(value, str):
value = [value]
self._jars = value
@property
def py_files(self) -> Optional[List[str]]:
"""Get the py_files of the component.
:return: The py_files of the component.
:rtype: Optional[List[str]]
"""
return self._py_files
@py_files.setter
def py_files(self, value):
"""Set the py_files of the component.
:param value: The py_files of the component.
:type value: Union[str, List[str]]
:return: No return
:rtype: None
"""
if isinstance(value, str):
value = [value]
self._py_files = value
def _to_dict(self) -> Dict:
result = super()._to_dict()
return result
def _to_rest_object(self):
result = super()._to_rest_object()
if "pyFiles" in result.properties.component_spec:
result.properties.component_spec["py_files"] = result.properties.component_spec.pop("pyFiles")
return result
@classmethod
def _from_rest_object_to_init_params(cls, obj) -> Dict:
if "py_files" in obj.properties.component_spec:
obj.properties.component_spec["pyFiles"] = obj.properties.component_spec.pop("py_files")
result = super()._from_rest_object_to_init_params(obj)
return result
