You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by gu...@apache.org on 2023/07/19 01:50:43 UTC
[spark] branch branch-3.5 updated: [SPARK-44264][ML][PYTHON] Refactoring TorchDistributor To Allow for Custom "run_training_on_file" Function Pointer
This is an automated email from the ASF dual-hosted git repository.
gurwls223 pushed a commit to branch branch-3.5
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-3.5 by this push:
new f4672b5fa5c [SPARK-44264][ML][PYTHON] Refactoring TorchDistributor To Allow for Custom "run_training_on_file" Function Pointer
f4672b5fa5c is described below
commit f4672b5fa5c229e9a76d45c620a07ec0ae1f81f4
Author: Mathew Jacob <ma...@databricks.com>
AuthorDate: Wed Jul 19 10:50:10 2023 +0900
[SPARK-44264][ML][PYTHON] Refactoring TorchDistributor To Allow for Custom "run_training_on_file" Function Pointer
### What Was Changed
We enable for a custom function pointer to be passed around the private functions that allow for distributed training of a function.
### Why Do We Need This Change
By abstracting the "run_training_on_pytorch_file" function to something that can be passed in, it allows for much easier creation of distributors that run on top of torch.distributed. Specifically, it makes it easy to implement distributed training of picklable functions in DeepspeedTorchDistributor. As mentioned, if there are accelerators that come out in the future built on top of torch.distributed, it will be very easy to support them in Spark. One can simply do the following:
1. Inherit from TorchDistributor and define a _run_training_on_pytorch_file function or equivalent for your class
2. When defining run(...), simply return _run() and pass in your custom _run_training_on_pytorch_file function in as the respective argument
### Any User-Facing Changes?
No.
### How Is This Tested?
The existing tests for TorchDistributor.
Closes #41973 from mathewjacob1002/distributed_func_support_prototype.
Lead-authored-by: Mathew Jacob <ma...@databricks.com>
Co-authored-by: Mathew Jacob <13...@users.noreply.github.com>
Signed-off-by: Hyukjin Kwon <gu...@apache.org>
(cherry picked from commit ee0e687b2bf50c6b8ab1b40b948e91ef03da42b8)
Signed-off-by: Hyukjin Kwon <gu...@apache.org>
---
python/pyspark/ml/torch/distributor.py | 117 ++++++++++++++++++++--
python/pyspark/ml/torch/tests/test_distributor.py | 9 +-
2 files changed, 113 insertions(+), 13 deletions(-)
diff --git a/python/pyspark/ml/torch/distributor.py b/python/pyspark/ml/torch/distributor.py
index 5b5d7c2288a..71bcde3b48e 100644
--- a/python/pyspark/ml/torch/distributor.py
+++ b/python/pyspark/ml/torch/distributor.py
@@ -515,10 +515,64 @@ class TorchDistributor(Distributor):
f"The {last_n_msg} included below: {task_output}"
)
+ @staticmethod
+ def _get_output_from_framework_wrapper(
+ framework_wrapper: Optional[Callable],
+ input_params: Dict,
+ train_object: Union[Callable, str],
+ run_pytorch_file_fn: Optional[Callable],
+ *args: Any,
+ **kwargs: Any,
+ ) -> Optional[Any]:
+ """
+ This function is meant to get the output from framework wrapper function by passing in the
+ correct arguments, depending on the type of train_object.
+
+ Parameters
+ ----------
+ framework_wrapper: Optional[Callable]
+ Function pointer that will be invoked. Can either be the function that runs distributed
+ training on files if train_object is a string. Otherwise, it will be the function that
+ runs distributed training for functions if the train_object is a Callable
+ input_params: Dict
+ A dictionary that maps parameter to arguments for the command to be created.
+ train_object: Union[Callable, str]
+ This input comes from the user. If the user inputs a string, then this means
+ it's a filepath. Otherwise, if the input is a function, then this means that
+ the user wants to run this function in a distributed manner.
+ run_pytorch_file_fn: Optional[Callable]
+ The function that will be used to run distributed training of a file;
+ mainly used for the distributed training using a function.
+ *args: Any
+ Extra arguments to be used by framework wrapper.
+ **kwargs: Any
+ Extra keyword args to be used. Not currently supported but kept for
+ future improvement.
+
+ Returns
+ -------
+ Optional[Any]
+ Returns the result of the framework_wrapper
+ """
+ if not framework_wrapper:
+ raise RuntimeError("`framework_wrapper` is not set. ...")
+ # The object to train is a file path, so framework_wrapper is some
+ # run_training_on_pytorch_file function.
+ if type(train_object) is str:
+ return framework_wrapper(input_params, train_object, *args, **kwargs)
+ else:
+ # We are doing training with a function, will call run_training_on_pytorch_function
+ if not run_pytorch_file_fn:
+ run_pytorch_file_fn = TorchDistributor._run_training_on_pytorch_file
+ return framework_wrapper(
+ input_params, train_object, run_pytorch_file_fn, *args, **kwargs
+ )
+
def _run_local_training(
self,
framework_wrapper_fn: Callable,
train_object: Union[Callable, str],
+ run_pytorch_file_fn: Optional[Callable],
*args: Any,
**kwargs: Any,
) -> Optional[Any]:
@@ -536,7 +590,14 @@ class TorchDistributor(Distributor):
os.environ[CUDA_VISIBLE_DEVICES] = ",".join(selected_gpus)
self.logger.info(f"Started local training with {self.num_processes} processes")
- output = framework_wrapper_fn(self.input_params, train_object, *args, **kwargs)
+ output = TorchDistributor._get_output_from_framework_wrapper(
+ framework_wrapper_fn,
+ self.input_params,
+ train_object,
+ run_pytorch_file_fn,
+ *args,
+ **kwargs,
+ )
self.logger.info(f"Finished local training with {self.num_processes} processes")
finally:
@@ -552,6 +613,7 @@ class TorchDistributor(Distributor):
self,
framework_wrapper_fn: Optional[Callable],
train_object: Union[Callable, str],
+ run_pytorch_file_fn: Optional[Callable],
input_dataframe: Optional["DataFrame"],
*args: Any,
**kwargs: Any,
@@ -656,7 +718,14 @@ class TorchDistributor(Distributor):
input_params["log_streaming_client"] = log_streaming_client
try:
with TorchDistributor._setup_spark_partition_data(iterator, schema_json):
- output = framework_wrapper_fn(input_params, train_object, *args, **kwargs)
+ output = TorchDistributor._get_output_from_framework_wrapper(
+ framework_wrapper_fn,
+ input_params,
+ train_object,
+ run_pytorch_file_fn,
+ *args,
+ **kwargs,
+ )
finally:
try:
LogStreamingClient._destroy()
@@ -686,6 +755,7 @@ class TorchDistributor(Distributor):
self,
framework_wrapper_fn: Callable,
train_object: Union[Callable, str],
+ run_pytorch_file_fn: Optional[Callable],
spark_dataframe: Optional["DataFrame"],
*args: Any,
**kwargs: Any,
@@ -702,7 +772,12 @@ class TorchDistributor(Distributor):
try:
spark_task_function = self._get_spark_task_function(
- framework_wrapper_fn, train_object, spark_dataframe, *args, **kwargs
+ framework_wrapper_fn,
+ train_object,
+ run_pytorch_file_fn,
+ spark_dataframe,
+ *args,
+ **kwargs,
)
self._check_encryption()
self.logger.info(
@@ -804,13 +879,21 @@ class TorchDistributor(Distributor):
@staticmethod
def _run_training_on_pytorch_function(
- input_params: Dict[str, Any], train_fn: Callable, *args: Any, **kwargs: Any
+ input_params: Dict[str, Any],
+ train_fn: Callable,
+ run_pytorch_file_fn: Optional[Callable],
+ *args: Any,
+ **kwargs: Any,
) -> Any:
+
+ if not run_pytorch_file_fn:
+ run_pytorch_file_fn = TorchDistributor._run_training_on_pytorch_file
+
with TorchDistributor._setup_files(train_fn, *args, **kwargs) as (
train_file_path,
output_file_path,
):
- TorchDistributor._run_training_on_pytorch_file(input_params, train_file_path)
+ run_pytorch_file_fn(input_params, train_file_path)
if not os.path.exists(output_file_path):
raise RuntimeError(
"TorchDistributor failed during training."
@@ -910,17 +993,28 @@ class TorchDistributor(Distributor):
train_object is a Callable with an expected output. Returns None if train_object is
a file.
"""
+ return self._run(
+ train_object, TorchDistributor._run_training_on_pytorch_file, *args, **kwargs
+ )
+
+ def _run(
+ self,
+ train_object: Union[Callable, str],
+ run_pytorch_file_fn: Callable,
+ *args: Any,
+ **kwargs: Any,
+ ) -> Optional[Any]:
if isinstance(train_object, str):
- framework_wrapper_fn = TorchDistributor._run_training_on_pytorch_file
+ framework_wrapper_fn = run_pytorch_file_fn
else:
- framework_wrapper_fn = (
- TorchDistributor._run_training_on_pytorch_function # type: ignore
- )
+ framework_wrapper_fn = TorchDistributor._run_training_on_pytorch_function
if self.local_mode:
- output = self._run_local_training(framework_wrapper_fn, train_object, *args, **kwargs)
+ output = self._run_local_training(
+ framework_wrapper_fn, train_object, run_pytorch_file_fn, *args, **kwargs
+ )
else:
output = self._run_distributed_training(
- framework_wrapper_fn, train_object, None, *args, **kwargs
+ framework_wrapper_fn, train_object, run_pytorch_file_fn, None, *args, **kwargs
)
return output
@@ -976,6 +1070,7 @@ class TorchDistributor(Distributor):
return self._run_distributed_training(
TorchDistributor._run_training_on_pytorch_function,
train_function,
+ TorchDistributor._run_training_on_pytorch_file,
spark_dataframe,
*args,
**kwargs,
diff --git a/python/pyspark/ml/torch/tests/test_distributor.py b/python/pyspark/ml/torch/tests/test_distributor.py
index 89c8893ea1b..364ed83f98d 100644
--- a/python/pyspark/ml/torch/tests/test_distributor.py
+++ b/python/pyspark/ml/torch/tests/test_distributor.py
@@ -383,7 +383,7 @@ class TorchDistributorLocalUnitTestsMixin:
)
self.assertEqual(
expected,
- dist._run_local_training(dist._run_training_on_pytorch_file, "train.py"),
+ dist._run_local_training(dist._run_training_on_pytorch_file, "train.py", None),
)
# cleanup
if cuda_env_var:
@@ -463,7 +463,12 @@ class TorchDistributorDistributedUnitTestsMixin:
)
self.assertEqual(
expected,
- dist._run_distributed_training(dist._run_training_on_pytorch_file, "...", None),
+ dist._run_distributed_training(
+ dist._run_training_on_pytorch_file,
+ "...",
+ TorchDistributor._run_training_on_pytorch_file,
+ None,
+ ),
)
def test_get_num_tasks_distributed(self) -> None:
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org