| # |
| # 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 |
| from contextlib import contextmanager |
| import collections |
| import logging |
| import math |
| import os |
| import random |
| import re |
| import shutil |
| import subprocess |
| import sys |
| import tempfile |
| import textwrap |
| import time |
| from typing import ( |
| Union, |
| Callable, |
| List, |
| Dict, |
| Optional, |
| Any, |
| Tuple, |
| Generator, |
| Iterator, |
| ) |
| |
| from pyspark import cloudpickle |
| from pyspark.resource.information import ResourceInformation |
| from pyspark.sql import DataFrame, SparkSession |
| from pyspark.taskcontext import BarrierTaskContext |
| from pyspark.ml.torch.log_communication import ( # type: ignore |
| LogStreamingClient, |
| LogStreamingServer, |
| ) |
| from pyspark.ml.dl_util import FunctionPickler |
| |
| |
| def _get_resources(session: SparkSession) -> Dict[str, ResourceInformation]: |
| resources: Dict[str, ResourceInformation] = {} |
| try: |
| resources = session.sparkContext.resources |
| except Exception: |
| resources = session._client._resources() # type: ignore[attr-defined] |
| return resources |
| |
| |
| def _get_conf(spark: SparkSession, key: str, default_value: str) -> str: |
| """Get the conf "key" from the given spark session, |
| or return the default value if the conf is not set. |
| |
| Parameters |
| ---------- |
| spark : :class:`SparkSession` |
| The :class:`SparkSession` for the distributor. |
| key : str |
| string for conf name |
| default_value : str |
| default value for the conf value for the given key |
| |
| Returns |
| ------- |
| str |
| Returns the string value that corresponds to the conf |
| """ |
| value = spark.conf.get(key, default_value) |
| assert value is not None |
| return value |
| |
| |
| # TODO(SPARK-41589): will move the functions and tests to an external file |
| # once we are in agreement about which functions should be in utils.py |
| def _get_conf_boolean(spark: SparkSession, key: str, default_value: str) -> bool: |
| value = _get_conf(spark=spark, key=key, default_value=default_value) |
| value = value.lower() |
| assert value in ["true", "false"] |
| return value == "true" |
| |
| |
| def _get_logger(name: str) -> logging.Logger: |
| """ |
| Gets a logger by name, or creates and configures it for the first time. |
| """ |
| logger = logging.getLogger(name) |
| logger.setLevel(logging.INFO) |
| # If the logger is configured, skip the configure |
| if not logger.handlers and not logging.getLogger().handlers: |
| handler = logging.StreamHandler(sys.stderr) |
| logger.addHandler(handler) |
| return logger |
| |
| |
| def _get_gpus_owned(context: Union[SparkSession, BarrierTaskContext]) -> List[str]: |
| """Gets the number of GPUs that Spark scheduled to the calling task. |
| |
| Parameters |
| ---------- |
| context : :class:`SparkSession` or :class:`BarrierTaskContext` |
| The :class:`SparkSession` or :class:`BarrierTaskContext` that has GPUs available. |
| |
| Returns |
| ------- |
| list |
| The correct mapping of addresses to workers. |
| |
| Raises |
| ------ |
| ValueError |
| Raised if the input addresses were not found. |
| """ |
| CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES" |
| pattern = re.compile("^[1-9][0-9]*|0$") |
| if isinstance(context, BarrierTaskContext): |
| addresses = context.resources()["gpu"].addresses |
| else: |
| addresses = _get_resources(context)["gpu"].addresses |
| |
| if any(not pattern.match(address) for address in addresses): |
| raise ValueError( |
| f"Found GPU addresses {addresses} which " |
| "are not all in the correct format " |
| "for CUDA_VISIBLE_DEVICES, which requires " |
| "integers with no zero padding." |
| ) |
| if CUDA_VISIBLE_DEVICES in os.environ: |
| gpu_indices = list(map(int, addresses)) |
| gpu_list = os.environ[CUDA_VISIBLE_DEVICES].split(",") |
| gpu_owned = [gpu_list[i] for i in gpu_indices] |
| return gpu_owned |
| return addresses |
| |
| |
| SPARK_PARTITION_ARROW_DATA_FILE = "SPARK_PARTITION_ARROW_DATA_FILE" |
| SPARK_DATAFRAME_SCHEMA_FILE = "SPARK_DATAFRAME_SCHEMA_FILE" |
| |
| |
| class Distributor: |
| """ |
| The parent class for TorchDistributor. This class shouldn't be instantiated directly. |
| """ |
| |
| def __init__( |
| self, |
| num_processes: int = 1, |
| local_mode: bool = True, |
| use_gpu: bool = True, |
| ssl_conf: Optional[str] = None, |
| ): |
| from pyspark.sql.utils import is_remote |
| |
| self.is_remote = is_remote() |
| self.spark = SparkSession.active() |
| |
| # indicate whether the server side is local mode |
| self.is_spark_local_master = False |
| # Refer to 'org.apache.spark.util.Utils#isLocalMaster' |
| master = _get_conf(self.spark, "spark.master", "") |
| if master == "local" or master.startswith("local["): |
| self.is_spark_local_master = True |
| |
| self.logger = _get_logger(self.__class__.__name__) |
| self.num_processes = num_processes |
| self.local_mode = local_mode |
| self.use_gpu = use_gpu |
| self.num_tasks = self._get_num_tasks() |
| self.ssl_conf = ssl_conf |
| |
| def _create_input_params(self) -> Dict[str, Any]: |
| input_params = self.__dict__.copy() |
| for unneeded_param in [ |
| "spark", |
| "ssl_conf", |
| "logger", |
| "is_remote", |
| "is_spark_local_master", |
| ]: |
| del input_params[unneeded_param] |
| return input_params |
| |
| def _get_num_tasks(self) -> int: |
| """ |
| Returns the number of Spark tasks to use for distributed training |
| |
| Returns |
| ------- |
| int |
| The number of Spark tasks to use for distributed training |
| |
| Raises |
| ------ |
| RuntimeError |
| Raised when the SparkConf was misconfigured. |
| """ |
| if self.use_gpu: |
| if not self.local_mode: |
| key = "spark.task.resource.gpu.amount" |
| task_gpu_amount = int(_get_conf(self.spark, key, "0")) |
| if task_gpu_amount < 1: |
| raise RuntimeError(f"'{key}' was unset, so gpu usage is unavailable.") |
| # TODO(SPARK-41916): Address situation when spark.task.resource.gpu.amount > 1 |
| return math.ceil(self.num_processes / task_gpu_amount) |
| else: |
| key = "spark.driver.resource.gpu.amount" |
| if "gpu" not in _get_resources(self.spark): |
| raise RuntimeError("GPUs were unable to be found on the driver.") |
| num_available_gpus = int(_get_conf(self.spark, key, "0")) |
| if num_available_gpus == 0: |
| raise RuntimeError("GPU resources were not configured properly on the driver.") |
| if self.num_processes > num_available_gpus: |
| self.logger.warning( |
| "'num_processes' cannot be set to a value greater than the number of " |
| f"available GPUs on the driver, which is {num_available_gpus}. " |
| "'num_processes' was reset to be equal to the number of available GPUs.", |
| ) |
| self.num_processes = num_available_gpus |
| return self.num_processes |
| |
| def _validate_input_params(self) -> None: |
| if self.num_processes <= 0: |
| raise ValueError("num_processes has to be a positive integer") |
| |
| def _check_encryption(self) -> None: |
| """Checks to see if the user requires encryption of data. |
| If required, throw an exception since we don't support that. |
| |
| Raises |
| ------ |
| RuntimeError |
| Thrown when the user requires ssl encryption or when the user initializes |
| the Distributor parent class. |
| """ |
| if not hasattr(self, "ssl_conf"): |
| raise RuntimeError( |
| "Distributor doesn't have this functionality. Use TorchDistributor instead." |
| ) |
| is_ssl_enabled = _get_conf_boolean(self.spark, "spark.ssl.enabled", "false") |
| ignore_ssl = _get_conf_boolean(self.spark, self.ssl_conf, "false") # type: ignore |
| if is_ssl_enabled: |
| name = self.__class__.__name__ |
| if ignore_ssl: |
| self.logger.warning( |
| textwrap.dedent( |
| f""" |
| This cluster has TLS encryption enabled; |
| however, {name} does not |
| support data encryption in transit. |
| The Spark configuration |
| '{self.ssl_conf}' has been set to |
| 'true' to override this |
| configuration and use {name} anyway. Please |
| note this will cause model |
| parameters and possibly training data to |
| be sent between nodes unencrypted. |
| """, |
| ) |
| ) |
| return |
| raise RuntimeError( |
| textwrap.dedent( |
| f""" |
| This cluster has TLS encryption enabled; |
| however, {name} does not support |
| data encryption in transit. To override |
| this configuration and use {name} |
| anyway, you may set '{self.ssl_conf}' |
| to 'true' in the Spark configuration. Please note this |
| will cause model parameters and possibly training |
| data to be sent between nodes unencrypted. |
| """ |
| ) |
| ) |
| |
| |
| class TorchDistributor(Distributor): |
| """ |
| A class to support distributed training on PyTorch and PyTorch Lightning using PySpark. |
| |
| .. versionadded:: 3.4.0 |
| |
| .. versionchanged:: 3.5.0 |
| Supports Spark Connect. |
| |
| Parameters |
| ---------- |
| num_processes : int, optional |
| An integer that determines how many different concurrent |
| tasks are allowed. We expect spark.task.gpus = 1 for GPU-enabled training. Default |
| should be 1; we don't want to invoke multiple cores/gpus without explicit mention. |
| local_mode : bool, optional |
| A boolean that determines whether we are using the driver |
| node for training. Default should be false; we don't want to invoke executors without |
| explicit mention. |
| use_gpu : bool, optional |
| A boolean that indicates whether or not we are doing training |
| on the GPU. Note that there are differences in how GPU-enabled code looks like and |
| how CPU-specific code looks like. |
| |
| Examples |
| -------- |
| Run PyTorch Training locally on GPU (using a PyTorch native function) |
| |
| >>> def train(learning_rate): |
| ... import torch.distributed |
| ... torch.distributed.init_process_group(backend="nccl") |
| ... # ... |
| ... torch.destroy_process_group() |
| ... return model # or anything else |
| ... |
| >>> distributor = TorchDistributor( |
| ... num_processes=2, |
| ... local_mode=True, |
| ... use_gpu=True) |
| >>> model = distributor.run(train, 1e-3) |
| |
| Run PyTorch Training on GPU (using a file with PyTorch code) |
| |
| >>> distributor = TorchDistributor( |
| ... num_processes=2, |
| ... local_mode=False, |
| ... use_gpu=True) |
| >>> distributor.run("/path/to/train.py", "--learning-rate=1e-3") |
| |
| Run PyTorch Lightning Training on GPU |
| |
| >>> num_proc = 2 |
| >>> def train(): |
| ... from pytorch_lightning import Trainer |
| ... # ... |
| ... # required to set devices = 1 and num_nodes = num_processes for multi node |
| ... # required to set devices = num_processes and num_nodes = 1 for single node multi GPU |
| ... trainer = Trainer(accelerator="gpu", devices=1, num_nodes=num_proc, strategy="ddp") |
| ... trainer.fit() |
| ... # ... |
| ... return trainer |
| ... |
| >>> distributor = TorchDistributor( |
| ... num_processes=num_proc, |
| ... local_mode=True, |
| ... use_gpu=True) |
| >>> trainer = distributor.run(train) |
| """ |
| |
| _PICKLED_FUNC_FILE = "func.pickle" |
| _TRAIN_FILE = "train.py" |
| _PICKLED_OUTPUT_FILE = "output.pickle" |
| _TORCH_SSL_CONF = "pytorch.spark.distributor.ignoreSsl" |
| |
| def __init__( |
| self, |
| num_processes: int = 1, |
| local_mode: bool = True, |
| use_gpu: bool = True, |
| _ssl_conf: str = _TORCH_SSL_CONF, |
| ): |
| """Initializes the distributor. |
| |
| Parameters |
| ---------- |
| num_processes : int, optional |
| An integer that determines how many different concurrent |
| tasks are allowed. We expect spark.task.gpus = 1 for GPU-enabled training. Default |
| should be 1; we don't want to invoke multiple cores/gpus without explicit mention. |
| local_mode : bool, optional |
| A boolean that determines whether we are using the driver |
| node for training. Default should be false; we don't want to invoke executors without |
| explicit mention. |
| use_gpu : bool, optional |
| A boolean that indicates whether or not we are doing training |
| on the GPU. Note that there are differences in how GPU-enabled code looks like and |
| how CPU-specific code looks like. |
| |
| Raises |
| ------ |
| ValueError |
| If any of the parameters are incorrect. |
| RuntimeError |
| If an active SparkSession is unavailable. |
| """ |
| super().__init__(num_processes, local_mode, use_gpu, ssl_conf=_ssl_conf) |
| self._validate_input_params() |
| self.input_params = self._create_input_params() |
| |
| @staticmethod |
| def _get_torchrun_args(local_mode: bool, num_processes: int) -> Tuple[List[Any], int]: |
| """ |
| Given the mode and the number of processes, create the arguments to be given to for torch |
| |
| Parameters |
| --------- |
| local_mode: bool |
| Whether or not we are running training locally or in a distributed fashion |
| |
| num_processes: int |
| The number of processes that we are going to use |
| |
| Returns |
| ------ |
| Tuple[List[Any], int] |
| A tuple containing a list of arguments to pass as pytorch args, |
| as well as the number of processes per node |
| """ |
| if local_mode: |
| torchrun_args = ["--standalone", "--nnodes=1"] |
| processes_per_node = num_processes |
| return torchrun_args, processes_per_node |
| |
| master_addr = os.environ["MASTER_ADDR"] |
| master_port = os.environ["MASTER_PORT"] |
| node_rank = os.environ["RANK"] |
| torchrun_args = [ |
| f"--nnodes={num_processes}", |
| f"--node_rank={node_rank}", |
| f"--rdzv_endpoint={master_addr}:{master_port}", |
| "--rdzv_id=0", # TODO: setup random ID that is gleaned from env variables |
| ] |
| processes_per_node = 1 |
| return torchrun_args, processes_per_node |
| |
| @staticmethod |
| def _create_torchrun_command( |
| input_params: Dict[str, Any], path_to_train_file: str, *args: Any |
| ) -> List[str]: |
| local_mode = input_params["local_mode"] |
| num_processes = input_params["num_processes"] |
| |
| torchrun_args, processes_per_node = TorchDistributor._get_torchrun_args( |
| local_mode=local_mode, num_processes=num_processes |
| ) |
| args_string = list(map(str, args)) # converting all args to strings |
| |
| return [ |
| sys.executable, |
| "-m", |
| "pyspark.ml.torch.torch_run_process_wrapper", |
| *torchrun_args, |
| f"--nproc_per_node={processes_per_node}", |
| path_to_train_file, |
| *args_string, |
| ] |
| |
| @staticmethod |
| def _execute_command( |
| cmd: List[str], |
| _prctl: bool = True, |
| redirect_to_stdout: bool = True, |
| log_streaming_client: Optional[LogStreamingClient] = None, |
| ) -> None: |
| _TAIL_LINES_TO_KEEP = 100 |
| |
| task = subprocess.Popen( |
| cmd, |
| stdout=subprocess.PIPE, |
| stderr=subprocess.STDOUT, |
| stdin=subprocess.PIPE, |
| env=os.environ, |
| ) |
| task.stdin.close() # type: ignore |
| tail: collections.deque = collections.deque(maxlen=_TAIL_LINES_TO_KEEP) |
| try: |
| for line in task.stdout: # type: ignore |
| decoded = line.decode() |
| tail.append(decoded) |
| if redirect_to_stdout: |
| if ( |
| log_streaming_client |
| and not log_streaming_client.failed |
| and ( |
| log_streaming_client.sock.getsockname()[0] |
| == log_streaming_client.sock.getpeername()[0] |
| ) |
| ): |
| # If log_streaming_client and log_stream_server are in the same |
| # node (typical case is spark local mode), |
| # server side will redirect the log to STDOUT, |
| # to avoid STDOUT outputs duplication, skip redirecting |
| # logs to STDOUT in client side. |
| pass |
| else: |
| sys.stdout.write(decoded) |
| if log_streaming_client: |
| log_streaming_client.send(decoded.rstrip()) |
| task.wait() |
| finally: |
| if task.poll() is None: |
| try: |
| task.terminate() # SIGTERM |
| time.sleep(0.5) |
| if task.poll() is None: |
| task.kill() # SIGKILL |
| except OSError: |
| pass |
| if task.returncode != os.EX_OK: |
| if len(tail) == _TAIL_LINES_TO_KEEP: |
| last_n_msg = f"last {_TAIL_LINES_TO_KEEP} lines of the task output are" |
| else: |
| last_n_msg = "task output is" |
| task_output = "".join(tail) |
| raise RuntimeError( |
| f"Command {cmd} failed with return code {task.returncode}. " |
| f"The {last_n_msg} included below: {task_output}" |
| ) |
| |
| @staticmethod |
| def _get_output_from_framework_wrapper( |
| framework_wrapper: Optional[Callable], |
| input_params: Dict, |
| train_object: Union[Callable, str], |
| run_pytorch_file_fn: Optional[Callable], |
| *args: Any, |
| **kwargs: Any, |
| ) -> Optional[Any]: |
| """ |
| This function is meant to get the output from framework wrapper function by passing in the |
| correct arguments, depending on the type of train_object. |
| |
| Parameters |
| ---------- |
| framework_wrapper: Optional[Callable] |
| Function pointer that will be invoked. Can either be the function that runs distributed |
| training on files if train_object is a string. Otherwise, it will be the function that |
| runs distributed training for functions if the train_object is a Callable |
| input_params: Dict |
| A dictionary that maps parameter to arguments for the command to be created. |
| train_object: Union[Callable, str] |
| This input comes from the user. If the user inputs a string, then this means |
| it's a filepath. Otherwise, if the input is a function, then this means that |
| the user wants to run this function in a distributed manner. |
| run_pytorch_file_fn: Optional[Callable] |
| The function that will be used to run distributed training of a file; |
| mainly used for the distributed training using a function. |
| *args: Any |
| Extra arguments to be used by framework wrapper. |
| **kwargs: Any |
| Extra keyword args to be used. Not currently supported but kept for |
| future improvement. |
| |
| Returns |
| ------- |
| Optional[Any] |
| Returns the result of the framework_wrapper |
| """ |
| if not framework_wrapper: |
| raise RuntimeError("`framework_wrapper` is not set. ...") |
| # The object to train is a file path, so framework_wrapper is some |
| # run_training_on_pytorch_file function. |
| if type(train_object) is str: |
| return framework_wrapper(input_params, train_object, *args, **kwargs) |
| else: |
| # We are doing training with a function, will call run_training_on_pytorch_function |
| if not run_pytorch_file_fn: |
| run_pytorch_file_fn = TorchDistributor._run_training_on_pytorch_file |
| return framework_wrapper( |
| input_params, train_object, run_pytorch_file_fn, *args, **kwargs |
| ) |
| |
| def _run_local_training( |
| self, |
| framework_wrapper_fn: Callable, |
| train_object: Union[Callable, str], |
| run_pytorch_file_fn: Optional[Callable], |
| *args: Any, |
| **kwargs: Any, |
| ) -> Optional[Any]: |
| CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES" |
| cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ |
| old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "") |
| try: |
| # Only replace the GPUs with 'SparkContext.resources' in legacy mode. |
| # In connect mode, this replacement is skipped since only GPUs on the client side |
| # can be used. |
| if self.use_gpu and not self.is_remote: |
| gpus_owned = _get_gpus_owned(self.spark) |
| random.seed(hash(train_object)) |
| selected_gpus = [str(e) for e in random.sample(gpus_owned, self.num_processes)] |
| os.environ[CUDA_VISIBLE_DEVICES] = ",".join(selected_gpus) |
| |
| self.logger.info(f"Started local training with {self.num_processes} processes") |
| output = TorchDistributor._get_output_from_framework_wrapper( |
| framework_wrapper_fn, |
| self.input_params, |
| train_object, |
| run_pytorch_file_fn, |
| *args, |
| **kwargs, |
| ) |
| self.logger.info(f"Finished local training with {self.num_processes} processes") |
| |
| finally: |
| if cuda_state_was_set: |
| os.environ[CUDA_VISIBLE_DEVICES] = old_cuda_visible_devices |
| else: |
| if CUDA_VISIBLE_DEVICES in os.environ: |
| del os.environ[CUDA_VISIBLE_DEVICES] |
| |
| return output |
| |
| def _get_spark_task_function( |
| self, |
| framework_wrapper_fn: Optional[Callable], |
| train_object: Union[Callable, str], |
| run_pytorch_file_fn: Optional[Callable], |
| input_dataframe: Optional["DataFrame"], |
| *args: Any, |
| **kwargs: Any, |
| ) -> Callable: |
| """Creates a spark task function that is used inside `mapPartitions`. |
| |
| Parameters |
| ---------- |
| framework_wrapper_fn : Optional[Callable] |
| The function that determines whether we are running training |
| on a PyTorch file or a PyTorch function. |
| train_object : Union[Callable, str] |
| The actual train function/file. |
| |
| Returns |
| ------- |
| Callable |
| The wrapped function ready for use with `mapPartitions` |
| """ |
| num_processes = self.num_processes |
| use_gpu = self.use_gpu |
| input_params = self.input_params |
| driver_address = self.driver_address |
| log_streaming_server_port = self.log_streaming_server_port |
| is_spark_local_master = self.is_spark_local_master |
| driver_owned_gpus: List[str] = [] |
| if is_spark_local_master and use_gpu: |
| driver_owned_gpus = _get_gpus_owned(self.spark) |
| |
| if input_dataframe is not None: |
| schema_json = input_dataframe.schema.jsonValue() |
| else: |
| schema_json = None |
| |
| # Spark task program |
| def wrapped_train_fn(iterator): # type: ignore[no-untyped-def] |
| import os |
| import pandas as pd |
| import pyarrow |
| from pyspark import BarrierTaskContext |
| |
| CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES" |
| |
| def get_free_port(address: str, context: "BarrierTaskContext") -> int: |
| port = "" |
| if context.partitionId() == 0: |
| try: |
| import socket |
| |
| sock = socket.socket() |
| sock.bind((address, 0)) |
| port = sock.getsockname()[1] |
| except socket.error: |
| pass |
| available_port = context.allGather(str(port))[0] |
| if not available_port: |
| raise RuntimeError("Failed to find free port for distributed training.") |
| return int(available_port) |
| |
| def set_torch_config(context: "BarrierTaskContext") -> None: |
| addrs = [e.address.split(":")[0] for e in context.getTaskInfos()] |
| |
| os.environ["MASTER_ADDR"] = str(addrs[0]) |
| os.environ["MASTER_PORT"] = str(get_free_port(addrs[0], context)) |
| os.environ["WORLD_SIZE"] = str(num_processes) |
| os.environ["NODE_RANK"] = str(context.partitionId()) |
| os.environ["RANK"] = str(context.partitionId()) |
| |
| if context.partitionId() >= num_processes: |
| raise ValueError( |
| "TorchDistributor._train_on_dataframe requires setting num_processes " |
| "equal to input spark dataframe partition number." |
| ) |
| |
| if is_spark_local_master: |
| # distributed training on a local mode spark cluster |
| def set_gpus(context: "BarrierTaskContext") -> None: |
| if CUDA_VISIBLE_DEVICES in os.environ: |
| return |
| |
| gpu_owned = driver_owned_gpus[context.partitionId()] |
| os.environ[CUDA_VISIBLE_DEVICES] = gpu_owned |
| |
| else: |
| |
| def set_gpus(context: "BarrierTaskContext") -> None: |
| if CUDA_VISIBLE_DEVICES in os.environ: |
| return |
| |
| gpus_owned = _get_gpus_owned(context) |
| os.environ[CUDA_VISIBLE_DEVICES] = ",".join(gpus_owned) |
| |
| context = BarrierTaskContext.get() |
| |
| if use_gpu: |
| set_gpus(context) |
| else: |
| os.environ[CUDA_VISIBLE_DEVICES] = "" |
| set_torch_config(context) |
| |
| log_streaming_client = LogStreamingClient(driver_address, log_streaming_server_port) |
| input_params["log_streaming_client"] = log_streaming_client |
| try: |
| with TorchDistributor._setup_spark_partition_data(iterator, schema_json): |
| output = TorchDistributor._get_output_from_framework_wrapper( |
| framework_wrapper_fn, |
| input_params, |
| train_object, |
| run_pytorch_file_fn, |
| *args, |
| **kwargs, |
| ) |
| finally: |
| try: |
| LogStreamingClient._destroy() |
| except BaseException: |
| pass |
| |
| if context.partitionId() == 0: |
| output_bytes = cloudpickle.dumps(output) |
| output_size = len(output_bytes) |
| |
| # In Spark Connect, DataFrame.collect stacks rows to size |
| # 'spark.connect.grpc.arrow.maxBatchSize' (default 4MiB), |
| # here use 4KiB for each chunk, which mean each arrow batch |
| # may contain about 1000 chunks. |
| chunks = [] |
| chunk_size = 4096 |
| index = 0 |
| while index < output_size: |
| chunks.append(output_bytes[index : index + chunk_size]) |
| index += chunk_size |
| |
| yield pyarrow.RecordBatch.from_pandas(pd.DataFrame(data={"chunk": chunks})) |
| |
| return wrapped_train_fn |
| |
| def _run_distributed_training( |
| self, |
| framework_wrapper_fn: Callable, |
| train_object: Union[Callable, str], |
| run_pytorch_file_fn: Optional[Callable], |
| spark_dataframe: Optional["DataFrame"], |
| *args: Any, |
| **kwargs: Any, |
| ) -> Optional[Any]: |
| log_streaming_server = LogStreamingServer() |
| self.driver_address = _get_conf(self.spark, "spark.driver.host", "") |
| assert self.driver_address != "" |
| try: |
| log_streaming_server.start(spark_host_address=self.driver_address) |
| time.sleep(1) # wait for the server to start |
| self.log_streaming_server_port = log_streaming_server.port |
| except Exception as e: |
| # If starting log streaming server failed, we don't need to break |
| # the distributor training but emit a warning instead. |
| self.log_streaming_server_port = -1 |
| self.logger.warning( |
| "Start torch distributor log streaming server failed, " |
| "You cannot receive logs sent from distributor workers, ", |
| f"error: {repr(e)}.", |
| ) |
| |
| try: |
| spark_task_function = self._get_spark_task_function( |
| framework_wrapper_fn, |
| train_object, |
| run_pytorch_file_fn, |
| spark_dataframe, |
| *args, |
| **kwargs, |
| ) |
| self._check_encryption() |
| self.logger.info( |
| f"Started distributed training with {self.num_processes} executor processes" |
| ) |
| if spark_dataframe is not None: |
| input_df = spark_dataframe |
| else: |
| input_df = self.spark.range( |
| start=0, end=self.num_tasks, step=1, numPartitions=self.num_tasks |
| ) |
| rows = input_df.mapInArrow( |
| func=spark_task_function, schema="chunk binary", barrier=True |
| ).collect() |
| output_bytes = b"".join([row.chunk for row in rows]) |
| result = cloudpickle.loads(output_bytes) |
| finally: |
| log_streaming_server.shutdown() |
| self.logger.info( |
| f"Finished distributed training with {self.num_processes} executor processes" |
| ) |
| return result |
| |
| @staticmethod |
| def _run_training_on_pytorch_file( |
| input_params: Dict[str, Any], train_path: str, *args: Any, **kwargs: Any |
| ) -> None: |
| if kwargs: |
| raise ValueError("Running pytorch file does not support key-word type arguments.") |
| log_streaming_client = input_params.get("log_streaming_client", None) |
| training_command = TorchDistributor._create_torchrun_command( |
| input_params, train_path, *args |
| ) |
| TorchDistributor._execute_command( |
| training_command, log_streaming_client=log_streaming_client |
| ) |
| |
| @staticmethod |
| @contextmanager |
| def _setup_files( |
| train_fn: Callable, *args: Any, **kwargs: Any |
| ) -> Generator[Tuple[str, str], None, None]: |
| save_dir = TorchDistributor._create_save_dir() |
| pickle_file_path = FunctionPickler.pickle_fn_and_save( |
| train_fn, "", save_dir, *args, **kwargs |
| ) |
| output_file_path = os.path.join(save_dir, TorchDistributor._PICKLED_OUTPUT_FILE) |
| script_path = os.path.join(save_dir, TorchDistributor._TRAIN_FILE) |
| train_file_path = FunctionPickler.create_fn_run_script( |
| pickle_file_path, output_file_path, script_path |
| ) |
| try: |
| yield (train_file_path, output_file_path) |
| finally: |
| TorchDistributor._cleanup_files(save_dir) |
| |
| @staticmethod |
| @contextmanager |
| def _setup_spark_partition_data( |
| partition_data_iterator: Iterator[Any], input_schema_json: Dict[str, Any] |
| ) -> Iterator[Any]: |
| from pyspark.sql.pandas.serializers import ArrowStreamSerializer |
| from pyspark.core.files import SparkFiles |
| import json |
| |
| if input_schema_json is None: |
| yield |
| return |
| |
| # We need to temporarily write partition data into a temp dir, |
| # partition data might be huge, so we need to write it under |
| # configured `SPARK_LOCAL_DIRS`. |
| save_dir = TorchDistributor._create_save_dir(root_dir=SparkFiles.getRootDirectory()) |
| |
| try: |
| serializer = ArrowStreamSerializer() |
| arrow_file_path = os.path.join(save_dir, "data.arrow") |
| with open(arrow_file_path, "wb") as f: |
| serializer.dump_stream(partition_data_iterator, f) |
| if f.tell() == 0: |
| # Nothing is written to file, this partition is empty |
| raise ValueError( |
| "Empty Spark partition is not allowed in " |
| "TorchDistributor.train_on_dataframe." |
| ) |
| |
| schema_file_path = os.path.join(save_dir, "schema.json") |
| schema_json_string = json.dumps(input_schema_json) |
| |
| with open(schema_file_path, "w") as f: |
| f.write(schema_json_string) |
| |
| os.environ[SPARK_PARTITION_ARROW_DATA_FILE] = arrow_file_path |
| os.environ[SPARK_DATAFRAME_SCHEMA_FILE] = schema_file_path |
| yield |
| finally: |
| os.environ.pop(SPARK_PARTITION_ARROW_DATA_FILE) |
| os.environ.pop(SPARK_DATAFRAME_SCHEMA_FILE) |
| TorchDistributor._cleanup_files(save_dir) |
| |
| @staticmethod |
| def _run_training_on_pytorch_function( |
| input_params: Dict[str, Any], |
| train_fn: Callable, |
| run_pytorch_file_fn: Optional[Callable], |
| *args: Any, |
| **kwargs: Any, |
| ) -> Any: |
| if not run_pytorch_file_fn: |
| run_pytorch_file_fn = TorchDistributor._run_training_on_pytorch_file |
| |
| with TorchDistributor._setup_files(train_fn, *args, **kwargs) as ( |
| train_file_path, |
| output_file_path, |
| ): |
| run_pytorch_file_fn(input_params, train_file_path) |
| if not os.path.exists(output_file_path): |
| raise RuntimeError( |
| "TorchDistributor failed during training." |
| "View stdout logs for detailed error message." |
| ) |
| try: |
| output = FunctionPickler.get_fn_output(output_file_path) |
| except Exception as e: |
| raise RuntimeError( |
| "TorchDistributor failed due to a pickling error. " |
| "View stdout logs for detailed error message." |
| ) from e |
| return output |
| |
| @staticmethod |
| def _create_save_dir(root_dir: Optional[str] = None) -> str: |
| # TODO: need to do this in a safe way to avoid issues during concurrent runs |
| return tempfile.mkdtemp(dir=root_dir) |
| |
| @staticmethod |
| def _cleanup_files(save_dir: str) -> None: |
| shutil.rmtree(save_dir, ignore_errors=True) |
| |
| def run(self, train_object: Union[Callable, str], *args: Any, **kwargs: Any) -> Optional[Any]: |
| """Runs distributed training. |
| |
| Parameters |
| ---------- |
| train_object : callable object or str |
| Either a PyTorch function, PyTorch Lightning function, or the path to a python file |
| that launches distributed training. |
| args : |
| If train_object is a python function and not a path to a python file, args need |
| to be the input parameters to that function. It would look like |
| |
| >>> model = distributor.run(train, 1e-3, 64) |
| |
| where train is a function and 1e-3 and 64 are regular numeric inputs to the function. |
| |
| If train_object is a python file, then args would be the command-line arguments for |
| that python file which are all in the form of strings. An example would be |
| |
| >>> distributor.run("/path/to/train.py", "--learning-rate=1e-3", "--batch-size=64") |
| |
| where since the input is a path, all of the parameters are strings that can be |
| handled by argparse in that python file. |
| kwargs : |
| If train_object is a python function and not a path to a python file, kwargs need |
| to be the key-word input parameters to that function. It would look like |
| |
| >>> model = distributor.run(train, tol=1e-3, max_iter=64) |
| |
| where train is a function of 2 arguments `tol` and `max_iter`. |
| |
| If train_object is a python file, then you should not set kwargs arguments. |
| |
| Returns |
| ------- |
| Returns the output of train_object called with args inside spark rank 0 task if the |
| train_object is a Callable with an expected output. Returns None if train_object is |
| a file. |
| """ |
| return self._run( |
| train_object, TorchDistributor._run_training_on_pytorch_file, *args, **kwargs |
| ) |
| |
| def _run( |
| self, |
| train_object: Union[Callable, str], |
| run_pytorch_file_fn: Callable, |
| *args: Any, |
| **kwargs: Any, |
| ) -> Optional[Any]: |
| if isinstance(train_object, str): |
| framework_wrapper_fn = run_pytorch_file_fn |
| else: |
| framework_wrapper_fn = TorchDistributor._run_training_on_pytorch_function |
| if self.local_mode: |
| output = self._run_local_training( |
| framework_wrapper_fn, train_object, run_pytorch_file_fn, *args, **kwargs |
| ) |
| else: |
| output = self._run_distributed_training( |
| framework_wrapper_fn, train_object, run_pytorch_file_fn, None, *args, **kwargs |
| ) |
| return output |
| |
| def _train_on_dataframe( |
| self, |
| train_function: Callable, |
| spark_dataframe: "DataFrame", |
| *args: Any, |
| **kwargs: Any, |
| ) -> Any: |
| """ |
| Runs distributed training using provided Spark DataFrame as input data. |
| You should ensure the input Spark DataFrame have evenly distributed partitions, |
| and this method starts a barrier Spark job that each Spark task in the job |
| process one partition of the input Spark DataFrame. |
| |
| Parameters |
| ---------- |
| train_function : |
| Either a PyTorch function, PyTorch Lightning function that launches distributed |
| training. Note that inside the function, you can call |
| `pyspark.ml.torch.distributor.get_spark_partition_data_loader` API to get a torch |
| data loader, the data loader loads data from the corresponding partition of the |
| input Spark DataFrame. |
| spark_dataframe : |
| An input Spark DataFrame that can be used in PyTorch `train_function` function. |
| See `train_function` argument doc for details. |
| args : |
| `args` need to be the input parameters to `train_function` function. It would look like |
| |
| >>> model = distributor.run(train, 1e-3, 64) |
| |
| where train is a function and 1e-3 and 64 are regular numeric inputs to the function. |
| kwargs : |
| `kwargs` need to be the key-word input parameters to `train_function` function. |
| It would look like |
| |
| >>> model = distributor.run(train, tol=1e-3, max_iter=64) |
| |
| where train is a function of 2 arguments `tol` and `max_iter`. |
| |
| Returns |
| ------- |
| Returns the output of `train_function` called with args inside Spark rank 0 task. |
| """ |
| |
| if self.local_mode: |
| raise ValueError( |
| "TorchDistributor.train_on_dataframe requires setting " |
| "TorchDistributor.local_mode to False." |
| ) |
| |
| return self._run_distributed_training( |
| TorchDistributor._run_training_on_pytorch_function, |
| train_function, |
| TorchDistributor._run_training_on_pytorch_file, |
| spark_dataframe, |
| *args, |
| **kwargs, |
| ) |
| |
| |
| def _get_spark_partition_data_loader( |
| num_samples: int, batch_size: int, num_workers: int = 1, prefetch_factor: int = 2 |
| ) -> Any: |
| """ |
| This function must be called inside the `train_function` where `train_function` |
| is the input argument of `TorchDistributor.train_on_dataframe`. |
| The function returns a pytorch data loader that loads data from |
| the corresponding spark partition data. |
| |
| Parameters |
| ---------- |
| num_samples : |
| Number of samples to generate per epoch. If `num_samples` is less than the number of |
| rows in the spark partition, it generate the first `num_samples` rows of |
| the spark partition, if `num_samples` is greater than the number of |
| rows in the spark partition, then after the iterator loaded all rows from the partition, |
| it wraps round back to the first row. |
| batch_size: |
| How many samples per batch to load. |
| num_workers: |
| How many subprocesses to use for data loading. |
| 0 means that the data will be loaded in the main process. |
| prefetch_factor: |
| Number of batches loaded in advance by each worker |
| """ |
| from pyspark.sql.types import StructType |
| from pyspark.ml.torch.data import _SparkPartitionTorchDataset |
| from torch.utils.data import DataLoader |
| |
| arrow_file = os.environ[SPARK_PARTITION_ARROW_DATA_FILE] |
| schema_file = os.environ[SPARK_DATAFRAME_SCHEMA_FILE] |
| |
| with open(schema_file, "r") as fp: |
| schema = StructType.fromJson(json.load(fp)) |
| |
| dataset = _SparkPartitionTorchDataset(arrow_file, schema, num_samples) |
| |
| if num_workers > 0: |
| return DataLoader( |
| dataset, batch_size, num_workers=num_workers, prefetch_factor=prefetch_factor |
| ) |
| else: |
| # if num_workers is zero, we cannot set `prefetch_factor` otherwise |
| # torch will raise error. |
| return DataLoader(dataset, batch_size, num_workers=num_workers) |
