Intro torchrun in CLI

pyproject.toml

@@ -25,6 +25,7 @@ dependencies = [
     "pyarrow",
     "pydantic>=2",
     "cloudpickle",
+    "importlib_metadata",
 ]
 dynamic = ["version"]
 classifiers = [

sdgx/cli/main.py

@@ -7,7 +7,7 @@
 import click
 
 from sdgx.cachers.manager import CacherManager
-from sdgx.cli.utils import cli_wrapper
+from sdgx.cli.utils import cli_wrapper, torch_run_warpper
 from sdgx.data_connectors.manager import DataConnectorManager
 from sdgx.data_exporters.manager import DataExporterManager
 from sdgx.data_processors.manager import DataProcessorManager
@@ -17,6 +17,7 @@
 
 
 @click.command()
+@torch_run_warpper
 @click.option(
     "--save_dir",
     type=str,
@@ -226,6 +227,7 @@ def fit(
 
 
 @click.command()
+@torch_run_warpper
 @click.option(
     "--load_dir",
     type=str,

sdgx/cli/utils.py

@@ -1,9 +1,15 @@
+import copy
+import json
+import re
+import sys
 from functools import wraps
 
 import click
+import importlib_metadata
 
 from sdgx.cli.message import ExceptionMessage, NormalMessage
 from sdgx.log import LOG_TO_FILE, add_log_file_handler, logger
+from sdgx.utils import find_free_port
 
 
 def cli_wrapper(func):
@@ -26,3 +32,54 @@ def wrapper(json_output, log_to_file, *args, **kwargs):
             exit(0)
 
     return wrapper
+
+
+def load_entry_point(distribution, group, name):
+    dist_obj = importlib_metadata.distribution(distribution)
+    eps = [ep for ep in dist_obj.entry_points if ep.group == group and ep.name == name]
+    if not eps:
+        raise ImportError("Entry point %r not found" % ((group, name),))
+    return eps[0].load()
+
+
+def torch_run_warpper(func):
+    """
+    Experimental feature for native torchrun.
+
+    Alternatively, people can use `torchrun $(which sdgx)`
+
+    FIXME: This is not compatible with click.testing.CliRunner
+    """
+
+    @click.option("--torchrun", type=bool, default=False, help="Use `torchrun` to run cli.")
+    @click.option(
+        "--torchrun_kwargs",
+        type=str,
+        default="{}",
+        help="[Json String] torchrun kwargs.",
+    )
+    @wraps(func)
+    def wrapper(torchrun, torchrun_kwargs, *args, **kwargs):
+        if not torchrun:
+            func(*args, **kwargs)
+        else:
+            torchrun_kwargs = json.loads(torchrun_kwargs)
+            torchrun_kwargs.setdefault("master_port", find_free_port())
+            origin_args = copy.deepcopy(sys.argv)
+            sys.argv = [re.sub(r"(-script\.pyw|\.exe)?$", "", sys.argv[0])]
+            for k, v in torchrun_kwargs.items():
+                sys.argv.extend([f"--{k}", str(v)])
+            # Remove [--torchrun=true] and [--torchrun, true] from origin_args
+            if "--torchrun" in origin_args:
+                i = origin_args.index("--torchrun")
+                if i + 1 < len(origin_args) and origin_args[i + 1] == "true":
+                    origin_args.pop(i)
+                    origin_args.pop(i)
+
+            if "--torchrun=true" in origin_args:
+                origin_args.remove("--torchrun=true")
+
+            sys.argv.extend(origin_args)
+            sys.exit(load_entry_point("torch", "console_scripts", "torchrun")())
+
+    return wrapper

sdgx/utils.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
 
+import socket
 import threading
 import urllib.request
+from contextlib import closing
 from pathlib import Path
 
 import pandas as pd
@@ -13,7 +15,14 @@
 except ImportError:
     from functools import lru_cache as cache
 
-__all__ = ["download_demo_data", "get_demo_single_table", "cache", "Singleton"]
+__all__ = ["download_demo_data", "get_demo_single_table", "cache", "Singleton", "find_free_port"]
+
+
+def find_free_port():
+    with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
+        s.bind(("", 0))
+        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        return s.getsockname()[1]
 
 
 def download_demo_data(data_dir: str | Path = "./dataset") -> Path:

tests/cli/test_cli.py

@@ -40,7 +40,10 @@ def test_list_extension_api(command, json_output):
 
 @pytest.mark.parametrize("model", ["CTGAN"])
 @pytest.mark.parametrize("json_output", [True, False])
-def test_fit_save_load_sample(model, demo_single_table_path, cacher_kwargs, json_output, tmp_path):
+@pytest.mark.parametrize("torchrun", [False])
+def test_fit_save_load_sample(
+    model, demo_single_table_path, cacher_kwargs, json_output, torchrun, tmp_path
+):
     runner = CliRunner()
     save_dir = tmp_path / f"unittest-{model}"
     result = runner.invoke(
@@ -62,6 +65,8 @@ def test_fit_save_load_sample(model, demo_single_table_path, cacher_kwargs, json
             json.dumps({"cacher_kwargs": cacher_kwargs}),
             "--json_output",
             json_output,
+            "--torchrun",
+            torchrun,
         ],
     )
 
@@ -84,6 +89,8 @@ def test_fit_save_load_sample(model, demo_single_table_path, cacher_kwargs, json
             json_output,
             "--export_dst",
             export_dst.as_posix(),
+            "--torchrun",
+            torchrun,
         ],
     )