You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by GitBox <gi...@apache.org> on 2022/12/23 02:34:44 UTC
[GitHub] [spark] rithwik-db opened a new pull request, #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
rithwik-db opened a new pull request, #39188:
URL: https://github.com/apache/spark/pull/39188
<!--
Thanks for sending a pull request! Here are some tips for you:
1. If this is your first time, please read our contributor guidelines: https://spark.apache.org/contributing.html
2. Ensure you have added or run the appropriate tests for your PR: https://spark.apache.org/developer-tools.html
3. If the PR is unfinished, add '[WIP]' in your PR title, e.g., '[WIP][SPARK-XXXX] Your PR title ...'.
4. Be sure to keep the PR description updated to reflect all changes.
5. Please write your PR title to summarize what this PR proposes.
6. If possible, provide a concise example to reproduce the issue for a faster review.
7. If you want to add a new configuration, please read the guideline first for naming configurations in
'core/src/main/scala/org/apache/spark/internal/config/ConfigEntry.scala'.
8. If you want to add or modify an error type or message, please read the guideline first in
'core/src/main/resources/error/README.md'.
-->
### What changes were proposed in this pull request?
<!--
Please clarify what changes you are proposing. The purpose of this section is to outline the changes and how this PR fixes the issue.
If possible, please consider writing useful notes for better and faster reviews in your PR. See the examples below.
1. If you refactor some codes with changing classes, showing the class hierarchy will help reviewers.
2. If you fix some SQL features, you can provide some references of other DBMSes.
3. If there is design documentation, please add the link.
4. If there is a discussion in the mailing list, please add the link.
-->
This is an addition to https://github.com/apache/spark/pull/39146 to add support for single node training using PyTorch files. The users would follow the second workflow in the [design document](https://docs.google.com/document/d/1QPO1Ly8WteL6aIPvVcR7Xne9qVtJiB3fdrRn7NwBcpA/edit#heading=h.8yvw9xq428fh) to run training. I added some new utility functions as well as built on top of current functions. This is largely WIP so testing will be added very soon.
### Why are the changes needed?
<!--
Please clarify why the changes are needed. For instance,
1. If you propose a new API, clarify the use case for a new API.
2. If you fix a bug, you can clarify why it is a bug.
-->
Look at the [main ticket](https://issues.apache.org/jira/browse/SPARK-41589) for more details.
### Does this PR introduce _any_ user-facing change?
<!--
Note that it means *any* user-facing change including all aspects such as the documentation fix.
If yes, please clarify the previous behavior and the change this PR proposes - provide the console output, description and/or an example to show the behavior difference if possible.
If possible, please also clarify if this is a user-facing change compared to the released Spark versions or within the unreleased branches such as master.
If no, write 'No'.
-->
No.
### How was this patch tested?
<!--
If tests were added, say they were added here. Please make sure to add some test cases that check the changes thoroughly including negative and positive cases if possible.
If it was tested in a way different from regular unit tests, please clarify how you tested step by step, ideally copy and paste-able, so that other reviewers can test and check, and descendants can verify in the future.
If tests were not added, please describe why they were not added and/or why it was difficult to add.
If benchmark tests were added, please run the benchmarks in GitHub Actions for the consistent environment, and the instructions could accord to: https://spark.apache.org/developer-tools.html#github-workflow-benchmarks.
-->
Testing isn't done yet; tests will be added soon.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] rithwik-db commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
rithwik-db commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1066491891
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -85,6 +130,12 @@ def __init__(
self.num_tasks = self._get_num_tasks()
self.ssl_conf = None
+ def _create_input_params(self) -> dict[str, Any]:
Review Comment:
This is just to pass the input params (specifically `num_processes` and `local_mode`) for the `staticmethod`s to use. We will have to have `kwargs` support in a future PR in case the users want to have a `max_restarts` option.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] lu-wang-dl commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
lu-wang-dl commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1066481463
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
+ try:
+ if self.use_gpu:
+ if "gpu" not in self.sc.resources:
+ raise ValueError('"gpu" was not found in the context resources.')
+ gpus_owned = get_gpus_owned(self.sc.resources["gpu"].addresses)
+
+ if self.num_processes > len(gpus_owned):
+ raise ValueError(
+ f"""{self.num_processes} processes were requested
+ for local training with GPU training but only
+ {len(gpus_owned)} GPUs were available."""
+ )
+ selected_gpus = [str(e) for e in random.sample(gpus_owned, self.num_processes)]
Review Comment:
Sure. It is fine if we add a dterministic random seed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] harupy commented on pull request #39188: [SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
harupy commented on PR #39188:
URL: https://github.com/apache/spark/pull/39188#issuecomment-1381288242
@rithwik-db Left a few comments on strings.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WeichenXu123 commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
WeichenXu123 commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1059226466
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
+ try:
+ if self.use_gpu:
+ if "gpu" not in self.sc.resources:
+ raise ValueError('"gpu" was not found in the context resources.')
+ gpus_owned = get_gpus_owned(self.sc.resources["gpu"].addresses)
+
+ if self.num_processes > len(gpus_owned):
+ raise ValueError(
+ f"""{self.num_processes} processes were requested
+ for local training with GPU training but only
+ {len(gpus_owned)} GPUs were available."""
+ )
+ selected_gpus = [str(e) for e in random.sample(gpus_owned, self.num_processes)]
+ os.environ[CUDA_VISIBLE_DEVICES] = ",".join(selected_gpus)
+
+ output = framework_wrapper_fn(self.__dict__, train_fn, *args)
+ finally:
+ if cuda_state_was_set:
+ os.environ[CUDA_VISIBLE_DEVICES] = old_cuda_visible_devices
+ else:
+ if (os.getenv(CUDA_VISIBLE_DEVICES, "")) != "":
+ del os.environ[CUDA_VISIBLE_DEVICES]
+
+ return output
+
+ @staticmethod
+ def _run_training_on_pytorch_file(
+ input_params: dict[str, Any], train_path: str, *args: Any
+ ) -> None:
+ training_command = create_torchrun_command(input_params, train_path, *args)
+ execute_command(training_command)
Review Comment:
Then I have a question,
if we running code on cluster mode spark cluster, then the "torchrun" subprocess will be executed in remote spark worker, then its output is also on remote worker disk, how can driver side get the output ?
Note that you cannot control which remote worker the spark task is scheduled to run on. This is very inconvenient to user.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] rithwik-db commented on a diff in pull request #39188: [SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
rithwik-db commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1068924576
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -261,6 +312,130 @@ def __init__(
super().__init__(num_processes, local_mode, use_gpu)
self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
self._validate_input_params()
+ self.input_params = self._create_input_params()
+
+ @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"]
+
+ if local_mode:
+ torchrun_args = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ pass
+ # TODO(SPARK-41592): Handle distributed training
Review Comment:
`local_mode` will always be true for this PR, but the follow up [PR](https://github.com/apache/spark/pull/39267) addresses the `else` section.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] rithwik-db commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
rithwik-db commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1065292487
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
+ try:
+ if self.use_gpu:
+ if "gpu" not in self.sc.resources:
+ raise ValueError('"gpu" was not found in the context resources.')
+ gpus_owned = get_gpus_owned(self.sc.resources["gpu"].addresses)
+
+ if self.num_processes > len(gpus_owned):
+ raise ValueError(
+ f"""{self.num_processes} processes were requested
+ for local training with GPU training but only
+ {len(gpus_owned)} GPUs were available."""
+ )
+ selected_gpus = [str(e) for e in random.sample(gpus_owned, self.num_processes)]
+ os.environ[CUDA_VISIBLE_DEVICES] = ",".join(selected_gpus)
+
+ output = framework_wrapper_fn(self.__dict__, train_fn, *args)
+ finally:
+ if cuda_state_was_set:
+ os.environ[CUDA_VISIBLE_DEVICES] = old_cuda_visible_devices
+ else:
+ if (os.getenv(CUDA_VISIBLE_DEVICES, "")) != "":
+ del os.environ[CUDA_VISIBLE_DEVICES]
+
+ return output
+
+ @staticmethod
+ def _run_training_on_pytorch_file(
+ input_params: dict[str, Any], train_path: str, *args: Any
+ ) -> None:
+ training_command = create_torchrun_command(input_params, train_path, *args)
+ execute_command(training_command)
Review Comment:
Should this just be in the official documentation or could it be inside some docstring in this file?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] HyukjinKwon commented on pull request #39188: [SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
HyukjinKwon commented on PR #39188:
URL: https://github.com/apache/spark/pull/39188#issuecomment-1379623516
Merged to master.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WeichenXu123 commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
WeichenXu123 commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1058964713
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
+ try:
+ if self.use_gpu:
+ if "gpu" not in self.sc.resources:
+ raise ValueError('"gpu" was not found in the context resources.')
+ gpus_owned = get_gpus_owned(self.sc.resources["gpu"].addresses)
+
+ if self.num_processes > len(gpus_owned):
+ raise ValueError(
+ f"""{self.num_processes} processes were requested
+ for local training with GPU training but only
+ {len(gpus_owned)} GPUs were available."""
+ )
+ selected_gpus = [str(e) for e in random.sample(gpus_owned, self.num_processes)]
+ os.environ[CUDA_VISIBLE_DEVICES] = ",".join(selected_gpus)
+
+ output = framework_wrapper_fn(self.__dict__, train_fn, *args)
+ finally:
+ if cuda_state_was_set:
+ os.environ[CUDA_VISIBLE_DEVICES] = old_cuda_visible_devices
+ else:
+ if (os.getenv(CUDA_VISIBLE_DEVICES, "")) != "":
Review Comment:
`if CUDA_VISIBLE_DEVICES in os.environ`
```suggestion
if (os.getenv(CUDA_VISIBLE_DEVICES, "")) != "":
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] rithwik-db commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
rithwik-db commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1059105070
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
+ try:
+ if self.use_gpu:
+ if "gpu" not in self.sc.resources:
+ raise ValueError('"gpu" was not found in the context resources.')
+ gpus_owned = get_gpus_owned(self.sc.resources["gpu"].addresses)
+
+ if self.num_processes > len(gpus_owned):
+ raise ValueError(
+ f"""{self.num_processes} processes were requested
+ for local training with GPU training but only
+ {len(gpus_owned)} GPUs were available."""
+ )
+ selected_gpus = [str(e) for e in random.sample(gpus_owned, self.num_processes)]
+ os.environ[CUDA_VISIBLE_DEVICES] = ",".join(selected_gpus)
+
+ output = framework_wrapper_fn(self.__dict__, train_fn, *args)
+ finally:
+ if cuda_state_was_set:
+ os.environ[CUDA_VISIBLE_DEVICES] = old_cuda_visible_devices
+ else:
+ if (os.getenv(CUDA_VISIBLE_DEVICES, "")) != "":
+ del os.environ[CUDA_VISIBLE_DEVICES]
+
+ return output
+
+ @staticmethod
+ def _run_training_on_pytorch_file(
+ input_params: dict[str, Any], train_path: str, *args: Any
+ ) -> None:
+ training_command = create_torchrun_command(input_params, train_path, *args)
+ execute_command(training_command)
Review Comment:
Yes, but to my understanding, that only happens when the input is a function. When the input is a file, the user code is expected to save the data themselves. There is no return value in these situations. Here's what I used for reference:
1. [PyTorch recommendation on how to create train.py files](https://pytorch.org/docs/stable/elastic/train_script.html)
2. [General (External) guide](https://lambdalabs.com/blog/multi-node-pytorch-distributed-training-guide)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] rithwik-db commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
rithwik-db commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1059228859
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
+ try:
+ if self.use_gpu:
+ if "gpu" not in self.sc.resources:
+ raise ValueError('"gpu" was not found in the context resources.')
+ gpus_owned = get_gpus_owned(self.sc.resources["gpu"].addresses)
+
+ if self.num_processes > len(gpus_owned):
+ raise ValueError(
+ f"""{self.num_processes} processes were requested
+ for local training with GPU training but only
+ {len(gpus_owned)} GPUs were available."""
+ )
+ selected_gpus = [str(e) for e in random.sample(gpus_owned, self.num_processes)]
+ os.environ[CUDA_VISIBLE_DEVICES] = ",".join(selected_gpus)
+
+ output = framework_wrapper_fn(self.__dict__, train_fn, *args)
+ finally:
+ if cuda_state_was_set:
+ os.environ[CUDA_VISIBLE_DEVICES] = old_cuda_visible_devices
+ else:
+ if (os.getenv(CUDA_VISIBLE_DEVICES, "")) != "":
+ del os.environ[CUDA_VISIBLE_DEVICES]
+
+ return output
+
+ @staticmethod
+ def _run_training_on_pytorch_file(
+ input_params: dict[str, Any], train_path: str, *args: Any
+ ) -> None:
+ training_command = create_torchrun_command(input_params, train_path, *args)
+ execute_command(training_command)
Review Comment:
Yes, we most likely need to force user to output the data to centralized storage.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] lu-wang-dl commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
lu-wang-dl commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1061987077
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
+ try:
+ if self.use_gpu:
+ if "gpu" not in self.sc.resources:
+ raise ValueError('"gpu" was not found in the context resources.')
+ gpus_owned = get_gpus_owned(self.sc.resources["gpu"].addresses)
+
+ if self.num_processes > len(gpus_owned):
+ raise ValueError(
+ f"""{self.num_processes} processes were requested
+ for local training with GPU training but only
+ {len(gpus_owned)} GPUs were available."""
+ )
+ selected_gpus = [str(e) for e in random.sample(gpus_owned, self.num_processes)]
Review Comment:
But this may cause more flaky issue if the user rerun the function. and make it more difficult to debug.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] rithwik-db commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
rithwik-db commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1059228859
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
+ try:
+ if self.use_gpu:
+ if "gpu" not in self.sc.resources:
+ raise ValueError('"gpu" was not found in the context resources.')
+ gpus_owned = get_gpus_owned(self.sc.resources["gpu"].addresses)
+
+ if self.num_processes > len(gpus_owned):
+ raise ValueError(
+ f"""{self.num_processes} processes were requested
+ for local training with GPU training but only
+ {len(gpus_owned)} GPUs were available."""
+ )
+ selected_gpus = [str(e) for e in random.sample(gpus_owned, self.num_processes)]
+ os.environ[CUDA_VISIBLE_DEVICES] = ",".join(selected_gpus)
+
+ output = framework_wrapper_fn(self.__dict__, train_fn, *args)
+ finally:
+ if cuda_state_was_set:
+ os.environ[CUDA_VISIBLE_DEVICES] = old_cuda_visible_devices
+ else:
+ if (os.getenv(CUDA_VISIBLE_DEVICES, "")) != "":
+ del os.environ[CUDA_VISIBLE_DEVICES]
+
+ return output
+
+ @staticmethod
+ def _run_training_on_pytorch_file(
+ input_params: dict[str, Any], train_path: str, *args: Any
+ ) -> None:
+ training_command = create_torchrun_command(input_params, train_path, *args)
+ execute_command(training_command)
Review Comment:
To do (1), we would need to provide the users with an interface for that right?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WeichenXu123 commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
WeichenXu123 commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1058946457
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
Review Comment:
I think checking parent pid becomes 1 will be a better solution for parent process died event.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] rithwik-db commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
rithwik-db commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1065291219
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -261,6 +312,129 @@ def __init__(
super().__init__(num_processes, local_mode, use_gpu)
self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
self._validate_input_params()
+ self.input_params = self._create_input_params()
+
+ @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"]
+
+ if local_mode:
+ torchrun_args = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ pass
+ # TODO(SPARK-41592): Handle distributed training
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["python", "-m", "pyspark.ml.torch.distributor.torch_run_process_wrapper"]
Review Comment:
Is this how I would call the wrapper @WeichenXu123?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] lu-wang-dl commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
lu-wang-dl commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1061987077
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
+ try:
+ if self.use_gpu:
+ if "gpu" not in self.sc.resources:
+ raise ValueError('"gpu" was not found in the context resources.')
+ gpus_owned = get_gpus_owned(self.sc.resources["gpu"].addresses)
+
+ if self.num_processes > len(gpus_owned):
+ raise ValueError(
+ f"""{self.num_processes} processes were requested
+ for local training with GPU training but only
+ {len(gpus_owned)} GPUs were available."""
+ )
+ selected_gpus = [str(e) for e in random.sample(gpus_owned, self.num_processes)]
Review Comment:
But this may cause more flaky issue if the user rerun the function.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] harupy commented on a diff in pull request #39188: [SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
harupy commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1068923718
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -261,6 +312,130 @@ def __init__(
super().__init__(num_processes, local_mode, use_gpu)
self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
self._validate_input_params()
+ self.input_params = self._create_input_params()
+
+ @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"]
+
+ if local_mode:
+ torchrun_args = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ pass
+ # TODO(SPARK-41592): Handle distributed training
Review Comment:
Not an issue if `local_mode` is always True, but code like below leaves `x` undefined when `condition` is False and causes an error.
```
if condition:
x =
else:
...
# Do something with x
print(x)
```
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -261,6 +312,130 @@ def __init__(
super().__init__(num_processes, local_mode, use_gpu)
self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
self._validate_input_params()
+ self.input_params = self._create_input_params()
+
+ @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"]
+
+ if local_mode:
+ torchrun_args = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ pass
+ # TODO(SPARK-41592): Handle distributed training
Review Comment:
Not an issue if `local_mode` is always True, but code like below leaves `x` undefined when `condition` is False and causes an error.
```python
if condition:
x =
else:
...
# Do something with x
print(x)
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] harupy commented on a diff in pull request #39188: [SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
harupy commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1068921168
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -261,6 +312,130 @@ def __init__(
super().__init__(num_processes, local_mode, use_gpu)
self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
self._validate_input_params()
+ self.input_params = self._create_input_params()
+
+ @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"]
+
+ if local_mode:
+ torchrun_args = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ pass
+ # TODO(SPARK-41592): Handle distributed training
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ [sys.executable, "-m", "pyspark.ml.torch.distributor.torch_run_process_wrapper"]
+ + torchrun_args
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [path_to_train_file, *args_string]
Review Comment:
```python
[
sys.executable,
"-m",
"pyspark.ml.torch.distributor.torch_run_process_wrapper",
*torch_args,
f"--nproc_per_node={processes_per_node}"
path_to_train_file,
*args_string
]
```
This might be simpler than concatenating four lists.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] rithwik-db commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
rithwik-db commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1059103421
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
Review Comment:
Sure
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WeichenXu123 commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
WeichenXu123 commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1058961073
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
+ try:
+ if self.use_gpu:
+ if "gpu" not in self.sc.resources:
+ raise ValueError('"gpu" was not found in the context resources.')
+ gpus_owned = get_gpus_owned(self.sc.resources["gpu"].addresses)
+
+ if self.num_processes > len(gpus_owned):
+ raise ValueError(
+ f"""{self.num_processes} processes were requested
+ for local training with GPU training but only
+ {len(gpus_owned)} GPUs were available."""
+ )
+ selected_gpus = [str(e) for e in random.sample(gpus_owned, self.num_processes)]
+ os.environ[CUDA_VISIBLE_DEVICES] = ",".join(selected_gpus)
+
+ output = framework_wrapper_fn(self.__dict__, train_fn, *args)
+ finally:
+ if cuda_state_was_set:
+ os.environ[CUDA_VISIBLE_DEVICES] = old_cuda_visible_devices
+ else:
+ if (os.getenv(CUDA_VISIBLE_DEVICES, "")) != "":
+ del os.environ[CUDA_VISIBLE_DEVICES]
+
+ return output
+
+ @staticmethod
+ def _run_training_on_pytorch_file(
+ input_params: dict[str, Any], train_path: str, *args: Any
+ ) -> None:
+ training_command = create_torchrun_command(input_params, train_path, *args)
+ execute_command(training_command)
+
+ def run(self, train_fn: Union[Callable, str], *args: Any) -> Optional[Any]:
+ """Runs distributed training on PySpark.
+
+ Parameters
+ ----------
+ train_fn : Union[Callable, str]
+ Either a PyTorch/PyTorch Lightning training function or the path to a python file
+ that launches distributed training.
Review Comment:
Param name is not good.
It looks like a python function but it can also be a path to a python code file.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] lu-wang-dl commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
lu-wang-dl commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1061884522
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
+ try:
+ if self.use_gpu:
+ if "gpu" not in self.sc.resources:
+ raise ValueError('"gpu" was not found in the context resources.')
+ gpus_owned = get_gpus_owned(self.sc.resources["gpu"].addresses)
+
+ if self.num_processes > len(gpus_owned):
+ raise ValueError(
+ f"""{self.num_processes} processes were requested
+ for local training with GPU training but only
+ {len(gpus_owned)} GPUs were available."""
+ )
+ selected_gpus = [str(e) for e in random.sample(gpus_owned, self.num_processes)]
+ os.environ[CUDA_VISIBLE_DEVICES] = ",".join(selected_gpus)
+
+ output = framework_wrapper_fn(self.__dict__, train_fn, *args)
+ finally:
+ if cuda_state_was_set:
+ os.environ[CUDA_VISIBLE_DEVICES] = old_cuda_visible_devices
+ else:
+ if (os.getenv(CUDA_VISIBLE_DEVICES, "")) != "":
+ del os.environ[CUDA_VISIBLE_DEVICES]
+
+ return output
+
+ @staticmethod
+ def _run_training_on_pytorch_file(
+ input_params: dict[str, Any], train_path: str, *args: Any
+ ) -> None:
+ training_command = create_torchrun_command(input_params, train_path, *args)
+ execute_command(training_command)
+
+ def run(self, train_fn: Union[Callable, str], *args: Any) -> Optional[Any]:
+ """Runs distributed training on PySpark.
+
+ Parameters
+ ----------
+ train_fn : Union[Callable, str]
+ Either a PyTorch/PyTorch Lightning training function or the path to a python file
+ that launches distributed training.
Review Comment:
`train_object` is fine with me.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] harupy commented on a diff in pull request #39188: [SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
harupy commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1068903058
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -261,6 +312,130 @@ def __init__(
super().__init__(num_processes, local_mode, use_gpu)
self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
self._validate_input_params()
+ self.input_params = self._create_input_params()
+
+ @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"]
+
+ if local_mode:
+ torchrun_args = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ pass
+ # TODO(SPARK-41592): Handle distributed training
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ [sys.executable, "-m", "pyspark.ml.torch.distributor.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
+ ) -> None:
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ,
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_object: Union[Callable, str],
+ *args: 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:
+ if self.use_gpu:
+ gpus_owned = get_gpus_owned(self.sc)
+
+ if self.num_processes > len(gpus_owned):
+ raise ValueError(
+ f"""{self.num_processes} processes were requested
+ for local training with GPU training but only
+ {len(gpus_owned)} GPUs were available."""
+ )
Review Comment:
I think the error message here would look like:
```
2 processes ...
for local ...
3 GPUs were ...
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WeichenXu123 commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
WeichenXu123 commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1065525284
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
+ try:
+ if self.use_gpu:
+ if "gpu" not in self.sc.resources:
+ raise ValueError('"gpu" was not found in the context resources.')
+ gpus_owned = get_gpus_owned(self.sc.resources["gpu"].addresses)
+
+ if self.num_processes > len(gpus_owned):
+ raise ValueError(
+ f"""{self.num_processes} processes were requested
+ for local training with GPU training but only
+ {len(gpus_owned)} GPUs were available."""
+ )
+ selected_gpus = [str(e) for e in random.sample(gpus_owned, self.num_processes)]
Review Comment:
@lu-wang-dl
We can fix the randomizer seed.
e.g., we can generate the randomizer seed by the hash of the torch program code,
then using the randomizer, we generate the selected_gpus list.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] rithwik-db commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
rithwik-db commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1065292784
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
+ try:
+ if self.use_gpu:
+ if "gpu" not in self.sc.resources:
+ raise ValueError('"gpu" was not found in the context resources.')
+ gpus_owned = get_gpus_owned(self.sc.resources["gpu"].addresses)
+
+ if self.num_processes > len(gpus_owned):
+ raise ValueError(
+ f"""{self.num_processes} processes were requested
+ for local training with GPU training but only
+ {len(gpus_owned)} GPUs were available."""
+ )
+ selected_gpus = [str(e) for e in random.sample(gpus_owned, self.num_processes)]
Review Comment:
So should we choose just the first `self.num_processes` GPUs owned then? I'm not sure what a good alternative approach here would be...
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] lu-wang-dl commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
lu-wang-dl commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1066482096
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -85,6 +130,12 @@ def __init__(
self.num_tasks = self._get_num_tasks()
self.ssl_conf = None
+ def _create_input_params(self) -> dict[str, Any]:
Review Comment:
Why do we need this instead of passing a `kwargs`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WeichenXu123 commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
WeichenXu123 commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1065522091
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
+ try:
+ if self.use_gpu:
+ if "gpu" not in self.sc.resources:
+ raise ValueError('"gpu" was not found in the context resources.')
+ gpus_owned = get_gpus_owned(self.sc.resources["gpu"].addresses)
+
+ if self.num_processes > len(gpus_owned):
+ raise ValueError(
+ f"""{self.num_processes} processes were requested
+ for local training with GPU training but only
+ {len(gpus_owned)} GPUs were available."""
+ )
+ selected_gpus = [str(e) for e in random.sample(gpus_owned, self.num_processes)]
+ os.environ[CUDA_VISIBLE_DEVICES] = ",".join(selected_gpus)
+
+ output = framework_wrapper_fn(self.__dict__, train_fn, *args)
+ finally:
+ if cuda_state_was_set:
+ os.environ[CUDA_VISIBLE_DEVICES] = old_cuda_visible_devices
+ else:
+ if (os.getenv(CUDA_VISIBLE_DEVICES, "")) != "":
+ del os.environ[CUDA_VISIBLE_DEVICES]
+
+ return output
+
+ @staticmethod
+ def _run_training_on_pytorch_file(
+ input_params: dict[str, Any], train_path: str, *args: Any
+ ) -> None:
+ training_command = create_torchrun_command(input_params, train_path, *args)
+ execute_command(training_command)
Review Comment:
We should add it in official doc.
and, we should also provide a databricks doc page with guide that contains the usage of copying data to dbfs.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WeichenXu123 commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
WeichenXu123 commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1058960621
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
Review Comment:
For local mode runs, only supports providing a training python file ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] rithwik-db commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
rithwik-db commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1061941062
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
+ try:
+ if self.use_gpu:
+ if "gpu" not in self.sc.resources:
+ raise ValueError('"gpu" was not found in the context resources.')
+ gpus_owned = get_gpus_owned(self.sc.resources["gpu"].addresses)
+
+ if self.num_processes > len(gpus_owned):
+ raise ValueError(
+ f"""{self.num_processes} processes were requested
+ for local training with GPU training but only
+ {len(gpus_owned)} GPUs were available."""
+ )
+ selected_gpus = [str(e) for e in random.sample(gpus_owned, self.num_processes)]
+ os.environ[CUDA_VISIBLE_DEVICES] = ",".join(selected_gpus)
+
+ output = framework_wrapper_fn(self.__dict__, train_fn, *args)
+ finally:
+ if cuda_state_was_set:
+ os.environ[CUDA_VISIBLE_DEVICES] = old_cuda_visible_devices
+ else:
+ if (os.getenv(CUDA_VISIBLE_DEVICES, "")) != "":
+ del os.environ[CUDA_VISIBLE_DEVICES]
+
+ return output
+
+ @staticmethod
+ def _run_training_on_pytorch_file(
+ input_params: dict[str, Any], train_path: str, *args: Any
+ ) -> None:
+ training_command = create_torchrun_command(input_params, train_path, *args)
+ execute_command(training_command)
+
+ def run(self, train_fn: Union[Callable, str], *args: Any) -> Optional[Any]:
+ """Runs distributed training on PySpark.
+
+ Parameters
+ ----------
+ train_fn : Union[Callable, str]
+ Either a PyTorch/PyTorch Lightning training function or the path to a python file
+ that launches distributed training.
+ args : *args
+ The arguments for train_fn
+
+ Returns
+ -------
+ Union[None, Any]
+ Returns the output of train_fn(*args) if framework is "pytorch-lightning"
Review Comment:
The only ones that cannot return values is if you input a file into `PyTorchDistributor(...).run()`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] lu-wang-dl commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
lu-wang-dl commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1061979910
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
Review Comment:
Mark these as internal function: `_execute_command`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] rithwik-db commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
rithwik-db commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1059103355
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
+ try:
+ if self.use_gpu:
+ if "gpu" not in self.sc.resources:
+ raise ValueError('"gpu" was not found in the context resources.')
+ gpus_owned = get_gpus_owned(self.sc.resources["gpu"].addresses)
+
+ if self.num_processes > len(gpus_owned):
+ raise ValueError(
+ f"""{self.num_processes} processes were requested
+ for local training with GPU training but only
+ {len(gpus_owned)} GPUs were available."""
+ )
+ selected_gpus = [str(e) for e in random.sample(gpus_owned, self.num_processes)]
+ os.environ[CUDA_VISIBLE_DEVICES] = ",".join(selected_gpus)
+
+ output = framework_wrapper_fn(self.__dict__, train_fn, *args)
Review Comment:
This is actually updated on https://github.com/apache/spark/pull/39267/commits/175562515a9501cec1a67a293c14b2c36b49c91d but I will move those changes down here since that makes more sense.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] rithwik-db commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
rithwik-db commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1059102935
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
+ try:
+ if self.use_gpu:
+ if "gpu" not in self.sc.resources:
+ raise ValueError('"gpu" was not found in the context resources.')
+ gpus_owned = get_gpus_owned(self.sc.resources["gpu"].addresses)
+
+ if self.num_processes > len(gpus_owned):
+ raise ValueError(
+ f"""{self.num_processes} processes were requested
+ for local training with GPU training but only
+ {len(gpus_owned)} GPUs were available."""
+ )
+ selected_gpus = [str(e) for e in random.sample(gpus_owned, self.num_processes)]
Review Comment:
If multiple runs happen concurrently, less chance that they will impact the same GPUs.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WeichenXu123 commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
WeichenXu123 commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1061399015
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
+ try:
+ if self.use_gpu:
+ if "gpu" not in self.sc.resources:
+ raise ValueError('"gpu" was not found in the context resources.')
+ gpus_owned = get_gpus_owned(self.sc.resources["gpu"].addresses)
+
+ if self.num_processes > len(gpus_owned):
+ raise ValueError(
+ f"""{self.num_processes} processes were requested
+ for local training with GPU training but only
+ {len(gpus_owned)} GPUs were available."""
+ )
+ selected_gpus = [str(e) for e in random.sample(gpus_owned, self.num_processes)]
+ os.environ[CUDA_VISIBLE_DEVICES] = ",".join(selected_gpus)
+
+ output = framework_wrapper_fn(self.__dict__, train_fn, *args)
+ finally:
+ if cuda_state_was_set:
+ os.environ[CUDA_VISIBLE_DEVICES] = old_cuda_visible_devices
+ else:
+ if (os.getenv(CUDA_VISIBLE_DEVICES, "")) != "":
+ del os.environ[CUDA_VISIBLE_DEVICES]
+
+ return output
+
+ @staticmethod
+ def _run_training_on_pytorch_file(
+ input_params: dict[str, Any], train_path: str, *args: Any
+ ) -> None:
+ training_command = create_torchrun_command(input_params, train_path, *args)
+ execute_command(training_command)
Review Comment:
> I assumed that if we tell them to store it within hdfs (or dbfs) then the driver could see that as well.
The assumption is correct.
We should document this.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] rithwik-db commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
rithwik-db commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1059102735
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
+ try:
+ if self.use_gpu:
+ if "gpu" not in self.sc.resources:
+ raise ValueError('"gpu" was not found in the context resources.')
+ gpus_owned = get_gpus_owned(self.sc.resources["gpu"].addresses)
+
+ if self.num_processes > len(gpus_owned):
+ raise ValueError(
+ f"""{self.num_processes} processes were requested
+ for local training with GPU training but only
+ {len(gpus_owned)} GPUs were available."""
+ )
+ selected_gpus = [str(e) for e in random.sample(gpus_owned, self.num_processes)]
+ os.environ[CUDA_VISIBLE_DEVICES] = ",".join(selected_gpus)
+
+ output = framework_wrapper_fn(self.__dict__, train_fn, *args)
+ finally:
+ if cuda_state_was_set:
+ os.environ[CUDA_VISIBLE_DEVICES] = old_cuda_visible_devices
+ else:
+ if (os.getenv(CUDA_VISIBLE_DEVICES, "")) != "":
+ del os.environ[CUDA_VISIBLE_DEVICES]
+
+ return output
+
+ @staticmethod
+ def _run_training_on_pytorch_file(
+ input_params: dict[str, Any], train_path: str, *args: Any
+ ) -> None:
+ training_command = create_torchrun_command(input_params, train_path, *args)
+ execute_command(training_command)
+
+ def run(self, train_fn: Union[Callable, str], *args: Any) -> Optional[Any]:
+ """Runs distributed training on PySpark.
+
+ Parameters
+ ----------
+ train_fn : Union[Callable, str]
+ Either a PyTorch/PyTorch Lightning training function or the path to a python file
+ that launches distributed training.
Review Comment:
Would you suggest `train_object` or just `train`? I am also not entirely sure about what the run(...) API should look like. @lu-wang-dl any thoughts from your side?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] rithwik-db commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
rithwik-db commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1059107025
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
Review Comment:
I agree, but I am not sure how to do that. Since we are calling `torchrun` with `subprocess.Popen([torchrun,...]...)`, where do we have the thread checking for parent id?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] rithwik-db commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
rithwik-db commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1059228859
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
+ try:
+ if self.use_gpu:
+ if "gpu" not in self.sc.resources:
+ raise ValueError('"gpu" was not found in the context resources.')
+ gpus_owned = get_gpus_owned(self.sc.resources["gpu"].addresses)
+
+ if self.num_processes > len(gpus_owned):
+ raise ValueError(
+ f"""{self.num_processes} processes were requested
+ for local training with GPU training but only
+ {len(gpus_owned)} GPUs were available."""
+ )
+ selected_gpus = [str(e) for e in random.sample(gpus_owned, self.num_processes)]
+ os.environ[CUDA_VISIBLE_DEVICES] = ",".join(selected_gpus)
+
+ output = framework_wrapper_fn(self.__dict__, train_fn, *args)
+ finally:
+ if cuda_state_was_set:
+ os.environ[CUDA_VISIBLE_DEVICES] = old_cuda_visible_devices
+ else:
+ if (os.getenv(CUDA_VISIBLE_DEVICES, "")) != "":
+ del os.environ[CUDA_VISIBLE_DEVICES]
+
+ return output
+
+ @staticmethod
+ def _run_training_on_pytorch_file(
+ input_params: dict[str, Any], train_path: str, *args: Any
+ ) -> None:
+ training_command = create_torchrun_command(input_params, train_path, *args)
+ execute_command(training_command)
Review Comment:
If the user provides a function that outputs a model or something like:
```
def train():
...
return model
output = Distributor().run(train)
```
Then we are going to be doing (1) as suggested here https://issues.apache.org/jira/browse/SPARK-41775
But I am assuming the issue at hand is when the user saves their `state_dict` or something like that in their train() function. In that case, I assumed that if we tell them to store it within hdfs (or dbfs) then the driver could see that as well. Is that assumption wrong?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] harupy commented on a diff in pull request #39188: [SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
harupy commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1068903858
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -261,6 +312,130 @@ def __init__(
super().__init__(num_processes, local_mode, use_gpu)
self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
self._validate_input_params()
+ self.input_params = self._create_input_params()
+
+ @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"]
+
+ if local_mode:
+ torchrun_args = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ pass
+ # TODO(SPARK-41592): Handle distributed training
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ [sys.executable, "-m", "pyspark.ml.torch.distributor.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
+ ) -> None:
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ,
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_object: Union[Callable, str],
+ *args: 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:
+ if self.use_gpu:
+ gpus_owned = get_gpus_owned(self.sc)
+
+ if self.num_processes > len(gpus_owned):
+ raise ValueError(
+ f"""{self.num_processes} processes were requested
+ for local training with GPU training but only
+ {len(gpus_owned)} GPUs were available."""
+ )
Review Comment:
Triple-quoted strings in `_check_encryption` also need to be fixed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] harupy commented on a diff in pull request #39188: [SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
harupy commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1068903058
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -261,6 +312,130 @@ def __init__(
super().__init__(num_processes, local_mode, use_gpu)
self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
self._validate_input_params()
+ self.input_params = self._create_input_params()
+
+ @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"]
+
+ if local_mode:
+ torchrun_args = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ pass
+ # TODO(SPARK-41592): Handle distributed training
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ [sys.executable, "-m", "pyspark.ml.torch.distributor.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
+ ) -> None:
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ,
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_object: Union[Callable, str],
+ *args: 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:
+ if self.use_gpu:
+ gpus_owned = get_gpus_owned(self.sc)
+
+ if self.num_processes > len(gpus_owned):
+ raise ValueError(
+ f"""{self.num_processes} processes were requested
+ for local training with GPU training but only
+ {len(gpus_owned)} GPUs were available."""
+ )
Review Comment:
This block results in a message like:
```
2 processes ...
for local ...
3 GPUs were ...
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] lu-wang-dl commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
lu-wang-dl commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1061884189
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
+ try:
+ if self.use_gpu:
+ if "gpu" not in self.sc.resources:
+ raise ValueError('"gpu" was not found in the context resources.')
+ gpus_owned = get_gpus_owned(self.sc.resources["gpu"].addresses)
+
+ if self.num_processes > len(gpus_owned):
+ raise ValueError(
+ f"""{self.num_processes} processes were requested
+ for local training with GPU training but only
+ {len(gpus_owned)} GPUs were available."""
+ )
+ selected_gpus = [str(e) for e in random.sample(gpus_owned, self.num_processes)]
+ os.environ[CUDA_VISIBLE_DEVICES] = ",".join(selected_gpus)
+
+ output = framework_wrapper_fn(self.__dict__, train_fn, *args)
+ finally:
+ if cuda_state_was_set:
+ os.environ[CUDA_VISIBLE_DEVICES] = old_cuda_visible_devices
+ else:
+ if (os.getenv(CUDA_VISIBLE_DEVICES, "")) != "":
+ del os.environ[CUDA_VISIBLE_DEVICES]
+
+ return output
+
+ @staticmethod
+ def _run_training_on_pytorch_file(
+ input_params: dict[str, Any], train_path: str, *args: Any
+ ) -> None:
+ training_command = create_torchrun_command(input_params, train_path, *args)
+ execute_command(training_command)
+
+ def run(self, train_fn: Union[Callable, str], *args: Any) -> Optional[Any]:
+ """Runs distributed training on PySpark.
+
+ Parameters
+ ----------
+ train_fn : Union[Callable, str]
+ Either a PyTorch/PyTorch Lightning training function or the path to a python file
+ that launches distributed training.
+ args : *args
+ The arguments for train_fn
+
+ Returns
+ -------
+ Union[None, Any]
+ Returns the output of train_fn(*args) if framework is "pytorch-lightning"
Review Comment:
Why only pytorch-lightning can return values?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] rithwik-db commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
rithwik-db commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1065290944
##########
python/pyspark/ml/torch/torch_run_process_wrapper.py:
##########
@@ -0,0 +1,80 @@
+#
Review Comment:
I will need to add tests for this wrapper but I am not sure how I would unit test it... Would I put all the code inside a function called `main()` and just unit test that?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WeichenXu123 commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
WeichenXu123 commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1065525284
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
+ try:
+ if self.use_gpu:
+ if "gpu" not in self.sc.resources:
+ raise ValueError('"gpu" was not found in the context resources.')
+ gpus_owned = get_gpus_owned(self.sc.resources["gpu"].addresses)
+
+ if self.num_processes > len(gpus_owned):
+ raise ValueError(
+ f"""{self.num_processes} processes were requested
+ for local training with GPU training but only
+ {len(gpus_owned)} GPUs were available."""
+ )
+ selected_gpus = [str(e) for e in random.sample(gpus_owned, self.num_processes)]
Review Comment:
@lu-wang-dl
We can set deterministic randomizer seed.
e.g., we can generate the randomizer seed by the hash of the torch program code,
then using the randomizer, we generate the selected_gpus list.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] HyukjinKwon commented on a diff in pull request #39188: [SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
HyukjinKwon commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1066665964
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -261,6 +312,130 @@ def __init__(
super().__init__(num_processes, local_mode, use_gpu)
self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
self._validate_input_params()
+ self.input_params = self._create_input_params()
+
+ @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"]
+
+ if local_mode:
+ torchrun_args = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ pass
+ # TODO(SPARK-41592): Handle distributed training
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["python", "-m", "pyspark.ml.torch.distributor.torch_run_process_wrapper"]
Review Comment:
I think you should get the Python executable from `sys.executable`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WeichenXu123 commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
WeichenXu123 commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1058970269
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
Review Comment:
variable name `standalone` is not clear.
use `torchrun_options` instead
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WeichenXu123 commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
WeichenXu123 commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1058972117
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
+ try:
+ if self.use_gpu:
+ if "gpu" not in self.sc.resources:
+ raise ValueError('"gpu" was not found in the context resources.')
+ gpus_owned = get_gpus_owned(self.sc.resources["gpu"].addresses)
+
+ if self.num_processes > len(gpus_owned):
+ raise ValueError(
+ f"""{self.num_processes} processes were requested
+ for local training with GPU training but only
+ {len(gpus_owned)} GPUs were available."""
+ )
+ selected_gpus = [str(e) for e in random.sample(gpus_owned, self.num_processes)]
+ os.environ[CUDA_VISIBLE_DEVICES] = ",".join(selected_gpus)
+
+ output = framework_wrapper_fn(self.__dict__, train_fn, *args)
+ finally:
+ if cuda_state_was_set:
+ os.environ[CUDA_VISIBLE_DEVICES] = old_cuda_visible_devices
+ else:
+ if (os.getenv(CUDA_VISIBLE_DEVICES, "")) != "":
+ del os.environ[CUDA_VISIBLE_DEVICES]
+
+ return output
+
+ @staticmethod
+ def _run_training_on_pytorch_file(
+ input_params: dict[str, Any], train_path: str, *args: Any
+ ) -> None:
+ training_command = create_torchrun_command(input_params, train_path, *args)
+ execute_command(training_command)
Review Comment:
The torchrun children process should output a model (to disk or return a pickled model data?) and then the `Distributor.run` should load the model and return it ?
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
+ try:
+ if self.use_gpu:
+ if "gpu" not in self.sc.resources:
+ raise ValueError('"gpu" was not found in the context resources.')
+ gpus_owned = get_gpus_owned(self.sc.resources["gpu"].addresses)
+
+ if self.num_processes > len(gpus_owned):
+ raise ValueError(
+ f"""{self.num_processes} processes were requested
+ for local training with GPU training but only
+ {len(gpus_owned)} GPUs were available."""
+ )
+ selected_gpus = [str(e) for e in random.sample(gpus_owned, self.num_processes)]
+ os.environ[CUDA_VISIBLE_DEVICES] = ",".join(selected_gpus)
+
+ output = framework_wrapper_fn(self.__dict__, train_fn, *args)
+ finally:
+ if cuda_state_was_set:
+ os.environ[CUDA_VISIBLE_DEVICES] = old_cuda_visible_devices
+ else:
+ if (os.getenv(CUDA_VISIBLE_DEVICES, "")) != "":
+ del os.environ[CUDA_VISIBLE_DEVICES]
+
+ return output
+
+ @staticmethod
+ def _run_training_on_pytorch_file(
+ input_params: dict[str, Any], train_path: str, *args: Any
+ ) -> None:
+ training_command = create_torchrun_command(input_params, train_path, *args)
+ execute_command(training_command)
Review Comment:
The torchrun children process should output a model (to disk or return a pickled model data?) and then the `Distributor.run` should load back the model and return it ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WeichenXu123 commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
WeichenXu123 commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1058954735
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
Review Comment:
This param is meaningless, because we can only create one spark session or get active spark session.
So I suggest remove this param, and we can get spark session by:
`SparkSession.getActiveSession()`, if active spark session is None, raise error and ask user to create spark session first.
I don't recommend internally get spark session by `SparkSession.builder.getOrCreate()`, if no active spark session, this line will create a local mode spark session with default config, which might be not what user want.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WeichenXu123 commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
WeichenXu123 commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1058951199
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
Review Comment:
We shouldn't provide this param, but checking local mode or not by:
`sparkSession.sparkContext.master.startswith("local")`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] rithwik-db commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
rithwik-db commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1059101827
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
Review Comment:
For now, there is a different ticket https://issues.apache.org/jira/browse/SPARK-41775 that involves adding the support code for python functions as well. That ticket will really easily build on top of this as long as we have support for python files first.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] rithwik-db commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
rithwik-db commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1059228967
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
Review Comment:
Okay I will create a ticket for this.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WeichenXu123 commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
WeichenXu123 commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1065520925
##########
python/pyspark/ml/torch/torch_run_process_wrapper.py:
##########
@@ -0,0 +1,80 @@
+#
Review Comment:
No. You can add unit test for it. `torch_run_process_wrapper` is also a python module, in test code you can still import it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] lu-wang-dl commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
lu-wang-dl commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1061985290
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
Review Comment:
What does this error message mean?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] rithwik-db commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
rithwik-db commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1059228859
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
+ try:
+ if self.use_gpu:
+ if "gpu" not in self.sc.resources:
+ raise ValueError('"gpu" was not found in the context resources.')
+ gpus_owned = get_gpus_owned(self.sc.resources["gpu"].addresses)
+
+ if self.num_processes > len(gpus_owned):
+ raise ValueError(
+ f"""{self.num_processes} processes were requested
+ for local training with GPU training but only
+ {len(gpus_owned)} GPUs were available."""
+ )
+ selected_gpus = [str(e) for e in random.sample(gpus_owned, self.num_processes)]
+ os.environ[CUDA_VISIBLE_DEVICES] = ",".join(selected_gpus)
+
+ output = framework_wrapper_fn(self.__dict__, train_fn, *args)
+ finally:
+ if cuda_state_was_set:
+ os.environ[CUDA_VISIBLE_DEVICES] = old_cuda_visible_devices
+ else:
+ if (os.getenv(CUDA_VISIBLE_DEVICES, "")) != "":
+ del os.environ[CUDA_VISIBLE_DEVICES]
+
+ return output
+
+ @staticmethod
+ def _run_training_on_pytorch_file(
+ input_params: dict[str, Any], train_path: str, *args: Any
+ ) -> None:
+ training_command = create_torchrun_command(input_params, train_path, *args)
+ execute_command(training_command)
Review Comment:
If the user provides a function that outputs a model or something like:
```
def train():
...
return model
output = Distributor().run(train)
```
Then we are going to be doing (1) as suggested here https://issues.apache.org/jira/browse/SPARK-41775
But I am assuming the issue at hand is when the user saves their `state_dict` or something like that in their train() function. In that case, I assumed that if we tell them to store it within hdfs (or dbfs) then the driver could see that as well.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] HyukjinKwon commented on a diff in pull request #39188: [SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
HyukjinKwon commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1066665066
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -64,6 +73,42 @@ def get_conf_boolean(sc: SparkContext, key: str, default_value: str) -> bool:
)
+def get_gpus_owned(sc: SparkContext) -> list[str]:
+ """Gets the number of GPUs that Spark scheduled to the calling task.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext that has GPUs available.
+
+ Returns
+ -------
+ list[str]
Review Comment:
```suggestion
list
```
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -64,6 +73,42 @@ def get_conf_boolean(sc: SparkContext, key: str, default_value: str) -> bool:
)
+def get_gpus_owned(sc: SparkContext) -> list[str]:
+ """Gets the number of GPUs that Spark scheduled to the calling task.
+
+ Parameters
+ ----------
+ sc : SparkContext
Review Comment:
```suggestion
sc : :class:`SparkContext`
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WeichenXu123 commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
WeichenXu123 commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1058963498
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
+ try:
+ if self.use_gpu:
+ if "gpu" not in self.sc.resources:
+ raise ValueError('"gpu" was not found in the context resources.')
+ gpus_owned = get_gpus_owned(self.sc.resources["gpu"].addresses)
+
+ if self.num_processes > len(gpus_owned):
+ raise ValueError(
+ f"""{self.num_processes} processes were requested
+ for local training with GPU training but only
+ {len(gpus_owned)} GPUs were available."""
+ )
+ selected_gpus = [str(e) for e in random.sample(gpus_owned, self.num_processes)]
Review Comment:
Q: Why randomly pick GPUs ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] AmplabJenkins commented on pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
AmplabJenkins commented on PR #39188:
URL: https://github.com/apache/spark/pull/39188#issuecomment-1364560370
Can one of the admins verify this patch?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WeichenXu123 commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
WeichenXu123 commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1059227857
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
+ try:
+ if self.use_gpu:
+ if "gpu" not in self.sc.resources:
+ raise ValueError('"gpu" was not found in the context resources.')
+ gpus_owned = get_gpus_owned(self.sc.resources["gpu"].addresses)
+
+ if self.num_processes > len(gpus_owned):
+ raise ValueError(
+ f"""{self.num_processes} processes were requested
+ for local training with GPU training but only
+ {len(gpus_owned)} GPUs were available."""
+ )
+ selected_gpus = [str(e) for e in random.sample(gpus_owned, self.num_processes)]
+ os.environ[CUDA_VISIBLE_DEVICES] = ",".join(selected_gpus)
+
+ output = framework_wrapper_fn(self.__dict__, train_fn, *args)
+ finally:
+ if cuda_state_was_set:
+ os.environ[CUDA_VISIBLE_DEVICES] = old_cuda_visible_devices
+ else:
+ if (os.getenv(CUDA_VISIBLE_DEVICES, "")) != "":
+ del os.environ[CUDA_VISIBLE_DEVICES]
+
+ return output
+
+ @staticmethod
+ def _run_training_on_pytorch_file(
+ input_params: dict[str, Any], train_path: str, *args: Any
+ ) -> None:
+ training_command = create_torchrun_command(input_params, train_path, *args)
+ execute_command(training_command)
Review Comment:
We have 2 solutions for this:
1) Force user code output model to specified path, then Distributor code can pickle model data and collect it to driver by RDD.collect()
2) (databricks only) Force user code output model to specified path, then Distributor code copy model to a specified dbfs directory that is accessible from driver side.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] lu-wang-dl commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
lu-wang-dl commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1061991762
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
+ try:
+ if self.use_gpu:
+ if "gpu" not in self.sc.resources:
+ raise ValueError('"gpu" was not found in the context resources.')
+ gpus_owned = get_gpus_owned(self.sc.resources["gpu"].addresses)
+
+ if self.num_processes > len(gpus_owned):
+ raise ValueError(
+ f"""{self.num_processes} processes were requested
+ for local training with GPU training but only
+ {len(gpus_owned)} GPUs were available."""
+ )
+ selected_gpus = [str(e) for e in random.sample(gpus_owned, self.num_processes)]
+ os.environ[CUDA_VISIBLE_DEVICES] = ",".join(selected_gpus)
+
+ output = framework_wrapper_fn(self.__dict__, train_fn, *args)
+ finally:
+ if cuda_state_was_set:
+ os.environ[CUDA_VISIBLE_DEVICES] = old_cuda_visible_devices
+ else:
+ if (os.getenv(CUDA_VISIBLE_DEVICES, "")) != "":
+ del os.environ[CUDA_VISIBLE_DEVICES]
+
+ return output
+
+ @staticmethod
+ def _run_training_on_pytorch_file(
+ input_params: dict[str, Any], train_path: str, *args: Any
+ ) -> None:
+ training_command = create_torchrun_command(input_params, train_path, *args)
+ execute_command(training_command)
Review Comment:
+1
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] rithwik-db commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
rithwik-db commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1061992443
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
Review Comment:
This is just a temporary error that will be removed down the line because once `pytorch_lightning` is supported, we will not have framework_wrapper_fn be optional.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] rithwik-db commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
rithwik-db commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1061940175
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
+ try:
+ if self.use_gpu:
+ if "gpu" not in self.sc.resources:
+ raise ValueError('"gpu" was not found in the context resources.')
+ gpus_owned = get_gpus_owned(self.sc.resources["gpu"].addresses)
+
+ if self.num_processes > len(gpus_owned):
+ raise ValueError(
+ f"""{self.num_processes} processes were requested
+ for local training with GPU training but only
+ {len(gpus_owned)} GPUs were available."""
+ )
+ selected_gpus = [str(e) for e in random.sample(gpus_owned, self.num_processes)]
+ os.environ[CUDA_VISIBLE_DEVICES] = ",".join(selected_gpus)
+
+ output = framework_wrapper_fn(self.__dict__, train_fn, *args)
+ finally:
+ if cuda_state_was_set:
+ os.environ[CUDA_VISIBLE_DEVICES] = old_cuda_visible_devices
+ else:
+ if (os.getenv(CUDA_VISIBLE_DEVICES, "")) != "":
+ del os.environ[CUDA_VISIBLE_DEVICES]
+
+ return output
+
+ @staticmethod
+ def _run_training_on_pytorch_file(
+ input_params: dict[str, Any], train_path: str, *args: Any
+ ) -> None:
+ training_command = create_torchrun_command(input_params, train_path, *args)
+ execute_command(training_command)
+
+ def run(self, train_fn: Union[Callable, str], *args: Any) -> Optional[Any]:
+ """Runs distributed training on PySpark.
+
+ Parameters
+ ----------
+ train_fn : Union[Callable, str]
+ Either a PyTorch/PyTorch Lightning training function or the path to a python file
+ that launches distributed training.
+ args : *args
+ The arguments for train_fn
+
+ Returns
+ -------
+ Union[None, Any]
+ Returns the output of train_fn(*args) if framework is "pytorch-lightning"
Review Comment:
Sorry that was an error on my part.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] lu-wang-dl commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
lu-wang-dl commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1061769537
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
Review Comment:
`local_mode` does not mean to use spark local mode, right? It will indicate that the code is running locally without spark, am I right?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] lu-wang-dl commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
lu-wang-dl commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1061983056
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
+ try:
+ if self.use_gpu:
+ if "gpu" not in self.sc.resources:
+ raise ValueError('"gpu" was not found in the context resources.')
Review Comment:
If GPU is not available, could we just run on CPU instead of throwing an error?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] rithwik-db commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
rithwik-db commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1065290535
##########
python/pyspark/ml/torch/tests/test_distributor.py:
##########
@@ -133,6 +214,42 @@ def test_get_num_tasks_locally(self) -> None:
distributor = TorchDistributor(num_processes, True, True)
distributor.num_processes = 3
+ def test_get_gpus_owned_local(self) -> None:
+ addresses = ["0", "1", "2"]
+ self.assertEqual(get_gpus_owned(self.sc), addresses)
+
+ env_vars = {"CUDA_VISIBLE_DEVICES": "3,4,5"}
+ self.setup_env_vars(env_vars)
+ self.assertEqual(get_gpus_owned(self.sc), ["3", "4", "5"])
+ self.delete_env_vars(env_vars)
+
+ def test_local_training_succeeds(self) -> None:
Review Comment:
We have a baseline test here, but there will need to be an integration test down the line (https://issues.apache.org/jira/browse/SPARK-41777)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] rithwik-db commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
rithwik-db commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1066263072
##########
python/pyspark/ml/torch/tests/test_distributor.py:
##########
@@ -178,6 +301,29 @@ def test_get_num_tasks_distributed(self) -> None:
self.spark.sparkContext._conf.set("spark.task.resource.gpu.amount", "1")
+class TorchWrapperUnitTests(unittest.TestCase):
Review Comment:
Same as [this](https://github.com/apache/spark/pull/39188/files#r1065290535), we need to add a separate integration test.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] rithwik-db commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
rithwik-db commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1059097298
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
Review Comment:
@lu-wang-dl is this okay with you?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] lu-wang-dl commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
lu-wang-dl commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1061968170
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
Review Comment:
We need to explain more on this. Have no idea what `train.py` mean.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] rithwik-db commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
rithwik-db commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1061943616
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
Review Comment:
The temporary commit for how this will be done is https://github.com/rithwik-db/spark_pytorch_distributor/commit/39ad2b76301e710283b41008ba63b1a9c5025ebe
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] rithwik-db commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
rithwik-db commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1061991285
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
+ try:
+ if self.use_gpu:
+ if "gpu" not in self.sc.resources:
+ raise ValueError('"gpu" was not found in the context resources.')
Review Comment:
But we would need to inform the user that setting `use_gpu=True` lead to that issue. Alternatively, we could tell the user that gpu was not found so we will switch to cpu through logging,
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
+ try:
+ if self.use_gpu:
+ if "gpu" not in self.sc.resources:
+ raise ValueError('"gpu" was not found in the context resources.')
Review Comment:
But we would need to inform the user that setting `use_gpu=True` lead to that issue. Alternatively, we could tell the user that gpu was not found so we will switch to cpu through logging.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] rithwik-db commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
rithwik-db commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1061991704
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
Review Comment:
Yeah, documentation will need to constantly evolve for this project.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WeichenXu123 commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
WeichenXu123 commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1065530249
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
+ try:
+ if self.use_gpu:
+ if "gpu" not in self.sc.resources:
+ raise ValueError('"gpu" was not found in the context resources.')
Review Comment:
I prefer raise error in this case, and tell user to set use_gpu=False because no GPU configured.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] rithwik-db commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
rithwik-db commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1066255412
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
+ try:
+ if self.use_gpu:
+ if "gpu" not in self.sc.resources:
+ raise ValueError('"gpu" was not found in the context resources.')
Review Comment:
We actually did this on this [line](https://github.com/apache/spark/blob/28a6efffc84555d20db4aa036d847cfad7d3e98a/python/pyspark/ml/torch/distributor.py#L164) (from the previous PR) so we don't need this check anymore.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] rithwik-db commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
rithwik-db commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1063016641
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
Review Comment:
Yes, sorry I misunderstood. `local_mode` essentially means run the code on the driver without using spark.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] harupy commented on a diff in pull request #39188: [SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
harupy commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1068903132
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -261,6 +312,130 @@ def __init__(
super().__init__(num_processes, local_mode, use_gpu)
self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
self._validate_input_params()
+ self.input_params = self._create_input_params()
+
+ @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"]
+
+ if local_mode:
+ torchrun_args = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ pass
+ # TODO(SPARK-41592): Handle distributed training
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ [sys.executable, "-m", "pyspark.ml.torch.distributor.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
+ ) -> None:
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ,
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}."
Review Comment:
I think this line is missing a space after `.`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] HyukjinKwon closed pull request #39188: [SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
HyukjinKwon closed pull request #39188: [SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
URL: https://github.com/apache/spark/pull/39188
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] harupy commented on a diff in pull request #39188: [SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
harupy commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1068903058
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -261,6 +312,130 @@ def __init__(
super().__init__(num_processes, local_mode, use_gpu)
self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
self._validate_input_params()
+ self.input_params = self._create_input_params()
+
+ @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"]
+
+ if local_mode:
+ torchrun_args = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ pass
+ # TODO(SPARK-41592): Handle distributed training
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ [sys.executable, "-m", "pyspark.ml.torch.distributor.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
+ ) -> None:
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ,
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_object: Union[Callable, str],
+ *args: 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:
+ if self.use_gpu:
+ gpus_owned = get_gpus_owned(self.sc)
+
+ if self.num_processes > len(gpus_owned):
+ raise ValueError(
+ f"""{self.num_processes} processes were requested
+ for local training with GPU training but only
+ {len(gpus_owned)} GPUs were available."""
+ )
Review Comment:
I think the error message would look like:
```
2 processes ...
for local ...
3 GPUs were ...
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] harupy commented on a diff in pull request #39188: [SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
harupy commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1068971084
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -261,6 +312,130 @@ def __init__(
super().__init__(num_processes, local_mode, use_gpu)
self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
self._validate_input_params()
+ self.input_params = self._create_input_params()
+
+ @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"]
+
+ if local_mode:
+ torchrun_args = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ pass
+ # TODO(SPARK-41592): Handle distributed training
Review Comment:
Makes sense!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] HyukjinKwon commented on a diff in pull request #39188: [SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
HyukjinKwon commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1066665271
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -64,6 +73,42 @@ def get_conf_boolean(sc: SparkContext, key: str, default_value: str) -> bool:
)
+def get_gpus_owned(sc: SparkContext) -> list[str]:
+ """Gets the number of GPUs that Spark scheduled to the calling task.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext that has GPUs available.
Review Comment:
```suggestion
The :class:`SparkContext` that has GPUs available.
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WeichenXu123 commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
WeichenXu123 commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1058969112
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
+ try:
+ if self.use_gpu:
+ if "gpu" not in self.sc.resources:
+ raise ValueError('"gpu" was not found in the context resources.')
+ gpus_owned = get_gpus_owned(self.sc.resources["gpu"].addresses)
+
+ if self.num_processes > len(gpus_owned):
+ raise ValueError(
+ f"""{self.num_processes} processes were requested
+ for local training with GPU training but only
+ {len(gpus_owned)} GPUs were available."""
+ )
+ selected_gpus = [str(e) for e in random.sample(gpus_owned, self.num_processes)]
+ os.environ[CUDA_VISIBLE_DEVICES] = ",".join(selected_gpus)
+
+ output = framework_wrapper_fn(self.__dict__, train_fn, *args)
Review Comment:
using `self.__dict__` is risky, we can only pick useful attributes and construct a param dict
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WeichenXu123 commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
WeichenXu123 commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1058964713
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
+ try:
+ if self.use_gpu:
+ if "gpu" not in self.sc.resources:
+ raise ValueError('"gpu" was not found in the context resources.')
+ gpus_owned = get_gpus_owned(self.sc.resources["gpu"].addresses)
+
+ if self.num_processes > len(gpus_owned):
+ raise ValueError(
+ f"""{self.num_processes} processes were requested
+ for local training with GPU training but only
+ {len(gpus_owned)} GPUs were available."""
+ )
+ selected_gpus = [str(e) for e in random.sample(gpus_owned, self.num_processes)]
+ os.environ[CUDA_VISIBLE_DEVICES] = ",".join(selected_gpus)
+
+ output = framework_wrapper_fn(self.__dict__, train_fn, *args)
+ finally:
+ if cuda_state_was_set:
+ os.environ[CUDA_VISIBLE_DEVICES] = old_cuda_visible_devices
+ else:
+ if (os.getenv(CUDA_VISIBLE_DEVICES, "")) != "":
Review Comment:
`if CUDA_VISIBLE_DEVICES in os.environ`
```suggestion
if CUDA_VISIBLE_DEVICES not in os.environ:
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WeichenXu123 commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
WeichenXu123 commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1059226466
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
+ )
+ 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:
+ sys.stdout.write(decoded)
+ 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}"
+ )
+
+
+class Distributor:
+ def __init__(
+ self,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ self.num_processes = num_processes
+ self.local_mode = local_mode
+ self.use_gpu = use_gpu
+ if spark:
+ self.spark = spark
+ else:
+ self.spark = SparkSession.builder.getOrCreate()
+ self.sc = self.spark.sparkContext
+ self.num_tasks = self._get_num_tasks()
+ self.ssl_conf = None
+
+ def _get_num_tasks(self) -> int:
+ """
+ Returns the number of Spark tasks to use for distributed training
+
+ Returns
+ -------
+ The number of Spark tasks to use for distributed training
+ """
+ if self.use_gpu:
+ key = "spark.task.resource.gpu.amount"
+ if self.sc.getConf().contains(key):
+ if gpu_amount_raw := self.sc.getConf().get(key): # mypy error??
+ task_gpu_amount = int(gpu_amount_raw)
+ else:
+ task_gpu_amount = 1 # for single node clusters
+ if task_gpu_amount < 1:
+ raise ValueError(
+ f"The Spark conf `{key}` has a value "
+ f"of {task_gpu_amount} but it "
+ "should not have a value less than 1."
+ )
+ return math.ceil(self.num_processes / task_gpu_amount)
+ return self.num_processes
+
+ def _validate_input_params(self) -> None:
+ if self.num_processes <= 0:
+ raise ValueError("num_proccesses has to be a positive integer")
+
+ def _check_encryption(self) -> None:
+ """Checks to see if the user requires encrpytion of data.
+ If required, throw an exception since we don't support that.
+
+ Raises
+ ------
+ NotImplementedError
+ Thrown when the user doesn't use PyTorchDistributor
+ Exception
+ Thrown when the user requires ssl encryption
+ """
+ if not "ssl_conf":
+ raise Exception(
+ "Distributor doesn't have this functionality. Use PyTorchDistributor instead."
+ )
+ is_ssl_enabled = get_conf_boolean(self.sc, "spark.ssl.enabled", "false")
+ ignore_ssl = get_conf_boolean(self.sc, self.ssl_conf, "false") # type: ignore
+ if is_ssl_enabled:
+ name = self.__class__.__name__
+ if ignore_ssl:
+ warnings.warn(
+ 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.
+ """,
+ RuntimeWarning,
+ )
+ return
+ raise Exception(
+ 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 PyTorchDistributor(Distributor):
+ """
+ A class to support distributed training on PyTorch and PyTorch Lightning using PySpark.
+
+ .. versionadded:: 3.4.0
+
+ 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()
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> distributor.run(train, 1e-3)
+
+ Run PyTorch Training on GPU (using a file with PyTorch code)
+
+ >>> distributor = Distributor(framework="pytorch",
+ num_processes=2,
+ local_mode=False,
+ use_gpu=True)
+ >>> distributor.run(/path/to/train.py, *args)
+
+ Run PyTorch Lightning Training
+ >>> def train():
+ >>> from pytorch_lightning import Trainer
+ >>> ...
+ >>> trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
+ >>> trainer.fit()
+ >>> ...
+ >>> return trainer
+ >>> distributor = Distributor(framework="pytorch-lightning",
+ num_processes=2,
+ local_mode=True,
+ use_gpu=True)
+ >>> trainer = distributor.run(train)
+ """
+
+ available_frameworks = ["pytorch", "pytorch-lightning"]
+ PICKLED_FUNC_FILE = "func.pickle"
+ TRAIN_FILE = "train.py"
+
+ def __init__(
+ self,
+ framework: str,
+ num_processes: int = 1,
+ local_mode: bool = True,
+ use_gpu: bool = True,
+ spark: Optional[SparkSession] = None,
+ ):
+ """Initializes the distributor.
+
+ Parameters
+ ----------
+ framework : str
+ A string indicating whether or not we are using PyTorch or PyTorch
+ Lightning. This could either be the string “pytorch” or ”pytorch-lightning”.
+ 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.
+ spark : Optional[SparkSession], optional
+ An optional parameter that allows users to pass in a custom SparkSession argument
+ with a custom conf, by default None
+
+ Raises
+ ------
+ ValueError
+ If any of the parameters are incorrect.
+ """
+ super().__init__(num_processes, local_mode, use_gpu, spark)
+ self.framework = framework
+ self.ssl_conf = "pytorch.spark.distributor.ignoreSsl" # type: ignore
+ self._validate_input_params()
+
+ def _validate_input_params(self) -> None:
+ """Validates input params
+
+ Raises
+ ------
+ ValueError
+ Thrown when user fails to provide correct input params
+ """
+ super()._validate_input_params()
+ if self.framework not in self.available_frameworks:
+ raise ValueError(
+ f"{self.framework} is not a valid framework."
+ f"Available frameworks: {self.available_frameworks}"
+ )
+
+ def _run_local_training(
+ self,
+ framework_wrapper_fn: Optional[Callable],
+ train_fn: Union[Callable, str],
+ *args: Any,
+ ) -> Optional[Any]:
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "")
+ cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ
+ if not framework_wrapper_fn:
+ raise RuntimeError("Unknown combination of parameters")
+ try:
+ if self.use_gpu:
+ if "gpu" not in self.sc.resources:
+ raise ValueError('"gpu" was not found in the context resources.')
+ gpus_owned = get_gpus_owned(self.sc.resources["gpu"].addresses)
+
+ if self.num_processes > len(gpus_owned):
+ raise ValueError(
+ f"""{self.num_processes} processes were requested
+ for local training with GPU training but only
+ {len(gpus_owned)} GPUs were available."""
+ )
+ selected_gpus = [str(e) for e in random.sample(gpus_owned, self.num_processes)]
+ os.environ[CUDA_VISIBLE_DEVICES] = ",".join(selected_gpus)
+
+ output = framework_wrapper_fn(self.__dict__, train_fn, *args)
+ finally:
+ if cuda_state_was_set:
+ os.environ[CUDA_VISIBLE_DEVICES] = old_cuda_visible_devices
+ else:
+ if (os.getenv(CUDA_VISIBLE_DEVICES, "")) != "":
+ del os.environ[CUDA_VISIBLE_DEVICES]
+
+ return output
+
+ @staticmethod
+ def _run_training_on_pytorch_file(
+ input_params: dict[str, Any], train_path: str, *args: Any
+ ) -> None:
+ training_command = create_torchrun_command(input_params, train_path, *args)
+ execute_command(training_command)
Review Comment:
Then I have a question,
if we running code on cluster mode spark cluster, then the "torchrun" subprocess will be executed in remote spark worker, then its output is also on remote worker disk, how can driver side get the output ?
Note that you cannot control which remote worker the spark task is scheduled to run on. This is very inconvenient to user. Shall we force user to output the data to centralized storage like dbfs ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WeichenXu123 commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
WeichenXu123 commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1059225843
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
+ preexec_fn=sigterm_on_parent_death,
Review Comment:
It's a good question. We can create a wrapping script to do that, the wrapping script creates `torchrun` subprocess, and spark code creates wrapping script subprocess.
We can do it later when we need to support all Unix system
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] lu-wang-dl commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
lu-wang-dl commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1061973010
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
+ """Returns the expected torchrun command based on the input.
+
+ Parameters
+ ----------
+ input_params : dict[str, Any]
+ The dictionary of the input parameters of the distributor. The most relevant params
+ are local_mode and num_processes.
+ train_path : str
+ The path to the (potentially autogenerated) train.py file
+ args: *args
+ The input arguments to the train.py file.
+
+ Returns
+ -------
+ str
+ The output torchrun command
+ """
+ local_mode = input_params["local_mode"]
+ num_processes = input_params["num_processes"]
+
+ if local_mode:
+ standalone = ["--standalone", "--nnodes=1"]
+ processes_per_node = num_processes
+ else:
+ master_addr, master_port = os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"]
+ node_rank = os.environ["RANK"]
+ standalone = [
+ 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
+
+ args_string = list(map(str, args)) # converting all args to strings
+
+ return (
+ ["torchrun"]
+ + standalone
+ + [f"--nproc_per_node={processes_per_node}"]
+ + [train_path, *args_string]
+ )
+
+
+def execute_command(cmd: list[str], _prctl: bool = True, redirect_to_stdout: bool = True) -> None:
+ """Runs a command in a new process, redirects stdout/stderr to stdout, and handles termination
+ of the command process.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ The users's input command split by the space characeter
+ _prctl : bool, optional
+ (test only) Use prctl to signal the command process upon parent death, by default True
+ redirect_to_stdout : bool, optional
+ whether to redirect all stdout/stderr logs to current stdout, by default True
+ TODO(SPARK-41593): Will update with additional parameters for logging to driver
+
+ Raises
+ ------
+ RuntimeError
+ If the child executed command fails
+ """
+ _TAIL_LINES_TO_KEEP = 100
+
+ def sigterm_on_parent_death() -> None:
+ """
+ Uses prctl to automatically send SIGTERM to the command process when its parent is dead.
+ This handles the case when the parent is a PySpark worker process.
+ If a user cancels the PySpark job, the worker process gets killed, regardless of
+ PySpark daemon and worker reuse settings.
+ NOTE: This doesn't work on windows or POSIX.
+ Future reference:
+ - MacOS/POSIX: https://stackoverflow.com/a/2035683
+ - Windows: https://stackoverflow.com/a/23587108
+ """
+ if _prctl:
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ # Set the parent process death signal of the command process to SIGTERM.
+ libc.prctl(1, signal.SIGTERM)
+ except OSError:
+ pass
+
+ task = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ stdin=subprocess.PIPE,
+ env=os.environ, # preserve the environ to support nested MPI jobs
Review Comment:
We don't need to mention MPI here.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] lu-wang-dl commented on a diff in pull request #39188: [WIP][SPARK-41591][PYTHON][ML] Training PyTorch Files on Single Node Multi GPU
Posted by GitBox <gi...@apache.org>.
lu-wang-dl commented on code in PR #39188:
URL: https://github.com/apache/spark/pull/39188#discussion_r1061979807
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -0,0 +1,491 @@
+#
+# 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 collections
+import ctypes
+import math
+import os
+import random
+import re
+import signal
+import sys
+import subprocess
+import time
+from typing import Union, Callable, Optional, Any
+import warnings
+
+from pyspark.sql import SparkSession
+from pyspark.context import SparkContext
+
+
+# Moved the util functions to this file for now
+# 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(sc: SparkContext, key: str, default_value: str) -> bool:
+ """Get the conf "key" from the given spark context,
+ or return the default value if the conf is not set.
+ This expects the conf value to be a boolean or string;
+ if the value is a string, this checks for all capitalization
+ patterns of "true" and "false" to match Scala.
+
+ Parameters
+ ----------
+ sc : SparkContext
+ The SparkContext for the distributor.
+ key : str
+ string for conf name
+ default_value : str
+ default value for the conf value for the given key
+
+ Returns
+ -------
+ bool
+ Returns the boolean value that corresponds to the conf
+
+ Raises
+ ------
+ Exception
+ Thrown when the conf value is not a boolean
+ """
+ val = sc.getConf().get(key, default_value)
+ lowercase_val = val.lower()
+ if lowercase_val == "true":
+ return True
+ if lowercase_val == "false":
+ return False
+ raise Exception(
+ "get_conf_boolean expected a boolean conf "
+ "value but found value of type {} "
+ "with value: {}".format(type(val), val)
+ )
+
+
+def get_gpus_owned(addresses: list[str]) -> list[str]:
+ """
+ Gets the number of GPUs that Spark scheduled to the calling task.
+ Returns:
+ The number of GPUs that Spark scheduled to the calling task.
+ """
+ CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+ pattern = re.compile("^[1-9][0-9]*|0$")
+ 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
+
+
+def create_torchrun_command(input_params: dict[str, Any], train_path: str, *args: Any) -> list[str]:
Review Comment:
Mark these as internal function: `_create_torchrun_command`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org