nurion-raydp 1.7.0__py3-none-any.whl

raydp/spark/dataset.py

@@ -0,0 +1,232 @@
+#
+# 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.
+import logging
+import uuid
+from typing import Callable, List, Optional, Union
+from dataclasses import dataclass
+
+import pandas as pd
+import pyarrow as pa
+import pyspark.sql as sql
+from pyspark.sql import SparkSession
+from pyspark.sql.dataframe import DataFrame
+from pyspark.sql.types import StructType
+from pyspark.sql.pandas.types import from_arrow_type
+from pyspark.storagelevel import StorageLevel
+import ray
+import ray.cross_language
+from ray.data import Dataset, from_arrow_refs
+from ray.types import ObjectRef
+from ray._private.client_mode_hook import client_mode_wrap
+
+from raydp.spark.ray_cluster_master import RAYDP_SPARK_MASTER_SUFFIX
+
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class PartitionObjectsOwner:
+    # Actor owner name
+    actor_name: str
+    # Function that set serialized parquet objects to actor owner state
+    # and return result of .remote() calling
+    set_reference_as_state: Callable[[ray.actor.ActorHandle, List[ObjectRef]], ObjectRef]
+
+
+def get_raydp_master_owner(
+    spark: Optional[SparkSession] = None,
+) -> PartitionObjectsOwner:
+    if spark is None:
+        spark = SparkSession.getActiveSession()
+    obj_holder_name = spark.sparkContext.appName + RAYDP_SPARK_MASTER_SUFFIX
+
+    def raydp_master_set_reference_as_state(
+        raydp_master_actor: ray.actor.ActorHandle, objects: List[ObjectRef]
+    ) -> ObjectRef:
+        return raydp_master_actor.add_objects.remote(uuid.uuid4(), objects)
+
+    return PartitionObjectsOwner(obj_holder_name, raydp_master_set_reference_as_state)
+
+
+@client_mode_wrap
+def _register_objects(records):
+    worker = ray.worker.global_worker
+    blocks: List[ray.ObjectRef] = []
+    block_sizes: List[int] = []
+    for obj_id, owner, num_record in records:
+        object_ref = ray.ObjectRef(obj_id)
+        # Register the ownership of the ObjectRef
+        worker.core_worker.deserialize_and_register_object_ref(
+            object_ref.binary(), ray.ObjectRef.nil(), owner, ""
+        )
+        blocks.append(object_ref)
+        block_sizes.append(num_record)
+    return blocks, block_sizes
+
+
+def _save_spark_df_to_object_store(
+    df: sql.DataFrame,
+    use_batch: bool = True,
+    owner: Union[PartitionObjectsOwner, None] = None,
+):
+    # call java function from python
+    jvm = df.sql_ctx.sparkSession.sparkContext._jvm
+    jdf = df._jdf
+    object_store_writer = jvm.org.apache.spark.sql.raydp.ObjectStoreWriter(jdf)
+    actor_owner_name = ""
+    if owner is not None:
+        actor_owner_name = owner.actor_name
+    records = object_store_writer.save(use_batch, actor_owner_name)
+
+    record_tuples = [
+        (record.objectId(), record.ownerAddress(), record.numRecords()) for record in records
+    ]
+    blocks, block_sizes = _register_objects(record_tuples)
+    logger.info(
+        f"after _register_objects, len(blocks): {len(blocks)}, len(block_sizes): {len(block_sizes)}"
+    )
+
+    if owner is not None:
+        actor_owner = ray.get_actor(actor_owner_name)
+        ray.get(owner.set_reference_as_state(actor_owner, blocks))
+
+    return blocks, block_sizes
+
+
+def spark_dataframe_to_ray_dataset(
+    df: sql.DataFrame,
+    parallelism: Optional[int] = None,
+    owner: Union[PartitionObjectsOwner, None] = None,
+):
+    num_part = df.rdd.getNumPartitions()
+    if parallelism is not None:
+        if parallelism != num_part:
+            df = df.repartition(parallelism)
+    blocks, _ = _save_spark_df_to_object_store(df, False, owner)
+    return from_arrow_refs(blocks)
+
+
+# This is an experimental API for now.
+# If you had any issue using it, welcome to report at our github.
+# This function WILL cache/persist the dataframe!
+def from_spark_recoverable(
+    df: sql.DataFrame,
+    storage_level: StorageLevel = StorageLevel.MEMORY_AND_DISK,
+    parallelism: Optional[int] = None,
+):
+    num_part = df.rdd.getNumPartitions()
+    if parallelism is not None:
+        if parallelism != num_part:
+            df = df.repartition(parallelism)
+    sc = df.sql_ctx.sparkSession.sparkContext
+    storage_level = sc._getJavaStorageLevel(storage_level)
+    object_store_writer = sc._jvm.org.apache.spark.sql.raydp.ObjectStoreWriter
+    object_ids = object_store_writer.fromSparkRDD(df._jdf, storage_level)
+    owner = object_store_writer.getAddress()
+    worker = ray.worker.global_worker
+    blocks = []
+    for object_id in object_ids:
+        object_ref = ray.ObjectRef(object_id)
+        # Register the ownership of the ObjectRef
+        worker.core_worker.deserialize_and_register_object_ref(
+            object_ref.binary(), ray.ObjectRef.nil(), owner, ""
+        )
+        blocks.append(object_ref)
+    return from_arrow_refs(blocks)
+
+
+def _convert_by_udf(
+    spark: sql.SparkSession,
+    blocks: List[ObjectRef],
+    locations: List[bytes],
+    schema: StructType,
+) -> DataFrame:
+    holder_name = spark.sparkContext.appName + RAYDP_SPARK_MASTER_SUFFIX
+    holder = ray.get_actor(holder_name)
+    df_id = uuid.uuid4()
+    ray.get(holder.add_objects.remote(df_id, blocks))
+    jvm = spark.sparkContext._jvm
+    object_store_reader = jvm.org.apache.spark.sql.raydp.ObjectStoreReader
+    # create the rdd then dataframe to utilize locality
+    jdf = object_store_reader.createRayObjectRefDF(spark._jsparkSession, locations)
+    current_namespace = ray.get_runtime_context().namespace
+    ray_address = ray.get(holder.get_ray_address.remote())
+    blocks_df = DataFrame(jdf, spark._wrapped if hasattr(spark, "_wrapped") else spark)
+
+    def _convert_blocks_to_dataframe(blocks):
+        # connect to ray
+        if not ray.is_initialized():
+            ray.init(
+                address=ray_address,
+                namespace=current_namespace,
+                logging_level=logging.WARN,
+            )
+        obj_holder = ray.get_actor(holder_name)
+        for block in blocks:
+            dfs = []
+            for idx in block["idx"]:
+                ref = ray.get(obj_holder.get_object.remote(df_id, idx))
+                data = ray.get(ref)
+                dfs.append(data.to_pandas())
+            yield pd.concat(dfs)
+
+    df = blocks_df.mapInPandas(_convert_blocks_to_dataframe, schema)
+    return df
+
+
+def _convert_by_rdd(
+    spark: sql.SparkSession, blocks: Dataset, locations: List[bytes], schema: StructType
+) -> DataFrame:
+    object_ids = [block.binary() for block in blocks]
+    schema_str = schema.json()
+    jvm = spark.sparkContext._jvm
+    # create rdd in java
+    rdd = jvm.org.apache.spark.rdd.RayDatasetRDD(spark._jsc, object_ids, locations)
+    # convert the rdd to dataframe
+    object_store_reader = jvm.org.apache.spark.sql.raydp.ObjectStoreReader
+    jdf = object_store_reader.RayDatasetToDataFrame(spark._jsparkSession, rdd, schema_str)
+    return DataFrame(jdf, spark._wrapped if hasattr(spark, "_wrapped") else spark)
+
+
+@client_mode_wrap
+def get_locations(blocks):
+    core_worker = ray.worker.global_worker.core_worker
+    return [core_worker.get_owner_address(block) for block in blocks]
+
+
+def ray_dataset_to_spark_dataframe(
+    spark: sql.SparkSession, arrow_schema, blocks: List[ObjectRef], locations=None
+) -> DataFrame:
+    locations = get_locations(blocks)
+    if hasattr(arrow_schema, "base_schema"):
+        arrow_schema = arrow_schema.base_schema
+    if not isinstance(arrow_schema, pa.lib.Schema):
+        raise RuntimeError(
+            f"Schema is {type(arrow_schema)}, required pyarrow.lib.Schema. \n"
+            f"to_spark does not support converting non-arrow ray datasets."
+        )
+    schema = StructType()
+    for field in arrow_schema:
+        schema.add(field.name, from_arrow_type(field.type), nullable=field.nullable)
+    # TODO how to branch on type of block?
+    sample = ray.get(blocks[0])
+    if isinstance(sample, bytes):
+        return _convert_by_rdd(spark, blocks, locations, schema)
+    elif isinstance(sample, pa.Table):
+        return _convert_by_udf(spark, blocks, locations, schema)
+    else:
+        raise RuntimeError("ray.to_spark only supports arrow type blocks")

raydp/spark/ray_cluster.py

@@ -0,0 +1,162 @@
+#
+# 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.
+#
+
+import glob
+import os
+import platform
+import pyspark
+from typing import Dict
+
+import ray
+from pyspark.sql.session import SparkSession
+
+from .ray_cluster_master import RAYDP_SPARK_MASTER_SUFFIX, RayDPSparkMaster
+
+DRIVER_CP_KEY = "spark.driver.extraClassPath"
+DRIVER_JAVA_OPTIONS_KEY = "spark.driver.extraJavaOptions"
+
+
+def _get_ray_job_id() -> str:
+    """Get the Ray job ID from environment or runtime context."""
+    # Try environment variable first
+    job_id = os.environ.get("RAY_JOB_ID")
+    if job_id:
+        return job_id
+    # Try runtime context
+    try:
+        ctx = ray.get_runtime_context()
+        return ctx.get_job_id()
+    except Exception:
+        return "local-job-id"
+
+
+class SparkCluster:
+    def __init__(
+        self,
+        app_name,
+        configs,
+        logging_level,
+    ):
+        self._app_name = app_name
+        self._configs = configs
+        self._logging_level = logging_level
+        # self._logger = logging.getLogger(__file__)
+        self._prepare_spark_configs()
+        self._setup_master(self._get_master_resources(self._configs))
+        self._spark_session: SparkSession = None
+
+    def _setup_master(self, resources: Dict[str, float]):
+        spark_master_name = self._app_name + RAYDP_SPARK_MASTER_SUFFIX
+
+        if resources:
+            num_cpu = 1
+            if "CPU" in resources:
+                num_cpu = resources["CPU"]
+                resources.pop("CPU", None)
+            self._spark_master_handle = RayDPSparkMaster.options(
+                name=spark_master_name,
+                num_cpus=num_cpu,
+                resources=resources,
+            ).remote(self._app_name, self._configs, logging_level=self._logging_level)
+        else:
+            self._spark_master_handle = RayDPSparkMaster.options(
+                name=spark_master_name,
+            ).remote(self._app_name, self._configs, logging_level=self._logging_level)
+
+        ray.get(self._spark_master_handle.start_up.remote(resources))
+
+    def _get_master_resources(self, configs: Dict[str, str]) -> Dict[str, float]:
+        resources = {}
+        spark_master_actor_resource_prefix = "spark.ray.master.actor.resource."
+
+        def get_master_actor_resource(
+            key_prefix: str, resource: Dict[str, float]
+        ) -> Dict[str, float]:
+            for key in configs:
+                if key.startswith(key_prefix):
+                    resource_name = key[len(key_prefix) :]
+                    resource[resource_name] = float(configs[key])
+            return resource
+
+        resources = get_master_actor_resource(spark_master_actor_resource_prefix, resources)
+
+        return resources
+
+    def get_cluster_url(self) -> str:
+        return ray.get(self._spark_master_handle.get_master_url.remote())
+
+    def _prepare_spark_configs(self):
+        if self._configs is None:
+            self._configs = {}
+        if platform.system() != "Darwin":
+            driver_node_ip = ray.util.get_node_ip_address()
+            if "spark.driver.host" not in self._configs:
+                self._configs["spark.driver.host"] = str(driver_node_ip)
+                self._configs["spark.driver.bindAddress"] = str(driver_node_ip)
+
+        raydp_cp = os.path.abspath(os.path.join(os.path.abspath(__file__), "../../jars/*"))
+        ray_cp = os.path.abspath(os.path.join(os.path.dirname(ray.__file__), "jars/*"))
+        spark_home = os.environ.get("SPARK_HOME", os.path.dirname(pyspark.__file__))
+        spark_jars_dir = os.path.abspath(os.path.join(spark_home, "jars/*"))
+
+        raydp_jars = glob.glob(raydp_cp)
+        driver_cp = ":".join(raydp_jars + [spark_jars_dir] + glob.glob(ray_cp))
+        if DRIVER_CP_KEY in self._configs:
+            self._configs[DRIVER_CP_KEY] += self._configs[DRIVER_CP_KEY] + ":" + driver_cp
+        else:
+            self._configs[DRIVER_CP_KEY] = driver_cp
+
+        extra_driver_options = f"-Dray.job.id={_get_ray_job_id()}"
+        if DRIVER_JAVA_OPTIONS_KEY in self._configs:
+            self._configs[DRIVER_JAVA_OPTIONS_KEY] += " " + extra_driver_options
+        else:
+            self._configs[DRIVER_JAVA_OPTIONS_KEY] = extra_driver_options
+
+        python_path_candidates = self._configs.get("spark.executorEnv.PYTHONPATH", "").split(":")
+        for k, v in os.environ.items():
+            if k == "PYTHONPATH":
+                python_path_candidates.append(v)
+            if k == "VIRTUAL_ENV":
+                python_path_candidates += glob.glob(f"{v}/lib/python*/site-packages")
+                self._configs["spark.pyspark.python"] = f"{v}/bin/python"
+        self._configs["spark.executorEnv.PYTHONPATH"] = ":".join(
+            [x for x in python_path_candidates if len(x) > 0]
+        )
+
+    def get_spark_session(self) -> SparkSession:
+        if self._spark_session is not None:
+            return self._spark_session
+        spark_builder = SparkSession.builder
+        for k, v in self._configs.items():
+            spark_builder.config(k, v)
+        spark_builder.enableHiveSupport()
+        self._spark_session = (
+            spark_builder.appName(self._app_name).master(self.get_cluster_url()).getOrCreate()
+        )
+
+        print(f"Spark UI: {self._spark_session.sparkContext.uiWebUrl}")
+        self._spark_session.sparkContext.setLogLevel(self._logging_level)
+        return self._spark_session
+
+    def stop(self, cleanup_data):
+        if self._spark_session is not None:
+            self._spark_session.stop()
+            self._spark_session = None
+        if self._spark_master_handle is not None:
+            self._spark_master_handle.stop.remote(cleanup_data)
+            if cleanup_data:
+                self._spark_master_handle = None

raydp/spark/ray_cluster_master.py

@@ -0,0 +1,102 @@
+#
+# 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.
+#
+
+import json
+import logging
+
+import ray
+import ray.cross_language
+from ray.util.scheduling_strategies import (
+    PlacementGroupSchedulingStrategy,
+)
+
+
+from .ray_pyworker import PyWorker
+
+RAYDP_SPARK_MASTER_SUFFIX = "_SPARK_MASTER"
+
+
+@ray.remote
+class RayDPSparkMaster:
+    def __init__(self, app_name, configs, logging_level: str):
+        self._logger = logging.getLogger(__file__)
+
+        self._app_name = app_name
+        self._ray_java_master = None
+        self._started_up = False
+        self._configs = configs
+        self._logging_level = logging_level
+        self._objects = {}
+
+    def start_up(self, resources=None):
+        if self._started_up:
+            self._logger.warning("The RayClusterMaster has started already. Do not call it twice")
+            return
+        ray_app_master_class = ray.cross_language.java_actor_class(
+            "org.apache.spark.deploy.raydp.RayAppMaster",
+            # {
+            #     "runtime_env": {
+            #         "java_executable": f"java -Dray.logging.level={self._logging_level} -cp {':'.join([p + '/*' for p in code_search_path()])}",
+            #     },
+            # },
+        )
+        self._logger.info(f"Start the RayClusterMaster with configs: {self._configs}")
+        self._ray_java_master = ray_app_master_class.options(resources=resources).remote()
+        self._logger.info("The RayClusterMaster has started")
+        self._started_up = True
+
+    def get_app_id(self) -> str:
+        assert self._started_up
+        return f"raydp-{self._ray_java_master._actor_id.hex()}"
+
+    def get_master_url(self) -> str:
+        assert self._started_up
+        url = ray.get(self._ray_java_master.getMasterUrl.remote())
+        self._logger.info(f"The master url is {url}")
+        return url
+
+    def create_pyworker(self, worker_id: str, node_id: str, env_vars: str) -> str:
+        self._logger.info(
+            f"Create a PyWorker with node_id: {node_id}, env_vars: {env_vars}, runtime_env: {ray.get_runtime_context().namespace}"
+        )
+        envs = json.loads(env_vars)
+        pg_name = f"raydp-executor-{self._app_name}-{worker_id}-pg"
+        # get placement group by name
+        pg = ray.util.get_placement_group(pg_name)
+        worker = PyWorker.options(
+            runtime_env={
+                "env_vars": envs,
+            },
+            max_concurrency=2,
+            scheduling_strategy=PlacementGroupSchedulingStrategy(pg),
+        ).remote()
+        ray.get(worker.heartbeat.remote())
+        return worker
+
+    def add_objects(self, timestamp, objects):
+        self._objects[timestamp] = objects
+
+    def get_object(self, timestamp, idx):
+        return self._objects[timestamp][idx]
+
+    def get_ray_address(self):
+        return ray.worker.global_worker.node.address
+
+    def stop(self, cleanup_data):
+        self._started_up = False
+        if cleanup_data:
+            ray.actor.exit_actor()

raydp/spark/ray_pyworker.py

@@ -0,0 +1,135 @@
+#
+# 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.
+#
+
+import logging
+import numbers
+import os
+import select
+import socket
+from errno import EINTR
+from socket import AF_INET, SOCK_STREAM, SOMAXCONN
+
+import ray
+
+logger = logging.getLogger(__file__)
+
+
+def compute_real_exit_code(exit_code):
+    # SystemExit's code can be integer or string, but os._exit only accepts integers
+    if isinstance(exit_code, numbers.Integral):
+        return exit_code
+    else:
+        return 1
+
+
+@ray.remote
+class PyWorker:
+    def __init__(self):
+        logger.info("PyWorker is created")
+
+        self.listen_sock = socket.socket(AF_INET, SOCK_STREAM)
+        self.listen_sock.bind(("127.0.0.1", 0))
+        self.listen_sock.listen(max(1024, SOMAXCONN))
+        listen_host, self.listen_port = self.listen_sock.getsockname()
+
+    def heartbeat(self):
+        return f"{os.getpid()} is alive"
+
+    def get_port(self) -> int:
+        return self.listen_port
+
+    def start(self):
+        import time
+        # Most of the code is copied from PySpark's daemon.py
+
+        from pyspark.serializers import (
+            UTF8Deserializer,
+            write_int,
+            write_with_length,
+        )
+        from pyspark.worker import main as worker_main
+
+        logger.info(f"Starting PyWorker with pid: {os.getpid()}, listen_port={self.listen_port}")
+        while True:
+            try:
+                logger.info("Waiting for connection")
+                ready_fds = select.select([0, self.listen_sock], [], [], 1)[0]
+            except select.error as ex:
+                logger.error(f"select error: {ex}")
+                if ex[0] == EINTR:
+                    continue
+                else:
+                    raise
+
+            logger.info(f"ready_fds: {ready_fds}")
+            if self.listen_sock in ready_fds:
+                try:
+                    sock, _ = self.listen_sock.accept()
+                except OSError as e:
+                    logger.error(f"Failed to accept connection: {e}")
+                    if e.errno == EINTR:
+                        continue
+                    raise
+
+                try:
+                    logger.info("Connection accepted")
+                    # Acknowledge that the fork was successful
+                    outfile = sock.makefile(mode="wb")
+                    write_int(os.getpid(), outfile)
+                    outfile.flush()
+                    outfile.close()
+                    while True:
+                        buffer_size = int(os.environ.get("SPARK_BUFFER_SIZE", 65536))
+                        infile = os.fdopen(os.dup(sock.fileno()), "rb", buffer_size)
+                        outfile = os.fdopen(os.dup(sock.fileno()), "wb", buffer_size)
+                        client_secret = UTF8Deserializer().loads(infile)
+                        if os.environ["PYTHON_WORKER_FACTORY_SECRET"] == client_secret:
+                            write_with_length("ok".encode("utf-8"), outfile)
+                            outfile.flush()
+                        else:
+                            write_with_length("err".encode("utf-8"), outfile)
+                            outfile.flush()
+                            sock.close()
+                            return 1
+
+                        try:
+                            code = worker_main(infile, outfile)
+                            logger.info(f"normal exit code: {code}")
+                        except SystemExit as exc:
+                            code = compute_real_exit_code(exc.code)
+                        finally:
+                            try:
+                                outfile.flush()
+                            except Exception:
+                                pass
+                        # wait for closing
+                        logger.info(f"exit code: {code}")
+                        # logger.info("Waiting for closing")
+                        # try:
+                        #     while sock.recv(1024):
+                        #         pass
+                        # except Exception:
+                        #     pass
+                        logger.info("Closing. Waiting for next loop")
+                        break
+                except BaseException as e:
+                    logger.error(f"PyWorker failed with exception: {e}")
+                    return 1
+                # else:
+                #     return 0
+            else:
+                time.sleep(0.5)