Add --policy presets for preview/unsupported gating

Author: godnight10061

API.md

@@ -8,6 +8,8 @@ import torchruntime
 You can use the command line:
 `python -m torchruntime install <optional list of package names and versions>`
 
+CLI flags: `--policy <compat|stable|preview>`, `--preview`, `--no-unsupported`, `--uv`
+
 Or you can use the library:
 ```py
 torchruntime.install(["torch", "torchvision<0.20"])

README.md

@@ -34,6 +34,10 @@ On Windows CUDA, Linux ROCm (6.x+), and Linux XPU, this also installs the approp
 
 **Tip:** You can also add the `--uv` flag to install packages using [uv](https://docs.astral.sh/uv/) (instead of `pip`). For e.g. `python -m torchruntime install --uv`
 
+Build-selection options:
+- `--policy <name>`: `compat` (default), `stable`, `preview` (or `nightly`)
+- Overrides: `--preview`, `--no-unsupported`
+
 ### Step 2. Configure torch
 This should be run inside your program, to initialize the required environment variables (if any) for the variant of torch being used.
 

tests/test_policy_parsing.py

@@ -0,0 +1,81 @@
+import pytest
+from torchruntime.utils.args import parse_policy_args
+
+
+def test_default_policy():
+    args = ["pkg1"]
+    preview, unsupported, cleaned = parse_policy_args(args)
+    assert preview is False
+    assert unsupported is True
+    assert cleaned == ["pkg1"]
+
+
+def test_stable_policy():
+    args = ["--policy", "stable", "pkg1"]
+    preview, unsupported, cleaned = parse_policy_args(args)
+    assert preview is False
+    assert unsupported is False
+    assert cleaned == ["pkg1"]
+
+
+def test_nightly_policy():
+    args = ["--policy", "nightly"]
+    preview, unsupported, cleaned = parse_policy_args(args)
+    assert preview is True
+    assert unsupported is True
+    assert cleaned == []
+
+
+def test_preview_policy_alias():
+    args = ["--policy", "preview"]
+    preview, unsupported, cleaned = parse_policy_args(args)
+    assert preview is True
+    assert unsupported is True
+    assert cleaned == []
+
+
+def test_policy_equals_syntax():
+    args = ["--policy=stable", "pkg1"]
+    preview, unsupported, cleaned = parse_policy_args(args)
+    assert preview is False
+    assert unsupported is False
+    assert cleaned == ["pkg1"]
+
+
+def test_policy_override_preview():
+    # stable is p=F, u=F. --preview should make p=T
+    args = ["--policy", "stable", "--preview"]
+    preview, unsupported, cleaned = parse_policy_args(args)
+    assert preview is True
+    assert unsupported is False
+
+def test_policy_override_unsupported():
+    # nightly is p=T, u=T. --no-unsupported should make u=F
+    args = ["--policy", "nightly", "--no-unsupported"]
+    preview, unsupported, cleaned = parse_policy_args(args)
+    assert preview is True
+    assert unsupported is False
+
+
+def test_unknown_policy():
+    args = ["--policy", "nonexistent"]
+    with pytest.raises(ValueError, match="Unknown policy"):
+        parse_policy_args(args)
+
+
+def test_missing_policy_arg():
+    args = ["--policy"]
+    with pytest.raises(ValueError, match="--policy requires an argument"):
+        parse_policy_args(args)
+
+
+def test_mixed_args():
+    args = ["torch", "--preview", "--policy", "stable", "--uv"]
+    # stable: p=F, u=F
+    # --preview: p=T
+    # Result: p=T, u=F
+    # cleaned: ["torch", "--uv"]
+    preview, unsupported, cleaned = parse_policy_args(args)
+    assert preview is True
+    assert unsupported is False
+    assert cleaned == ["torch", "--uv"]

torchruntime/__main__.py

@@ -1,6 +1,7 @@
 from .installer import install
 from .utils.torch_test import test
 from .utils import info
+from .utils.args import parse_policy_args
 
 
 def print_usage(entry_command: str):
@@ -18,6 +19,7 @@ def print_usage(entry_command: str):
     {entry_command} install --uv
     {entry_command} install --preview
     {entry_command} install --no-unsupported
+    {entry_command} install --policy stable
     {entry_command} install torch==2.2.0 torchvision==0.17.0
     {entry_command} install --uv torch>=2.0.0 torchaudio
     {entry_command} install torch==2.1.* torchvision>=0.16.0 torchaudio==2.1.0
@@ -39,6 +41,7 @@ def print_usage(entry_command: str):
     --uv               Use uv instead of pip for installation
     --preview          Allow preview builds (e.g. ROCm 6.4)
     --no-unsupported   Forbid EOL/unsupported builds (e.g. DirectML / IPEX / Torch 1.x)
+    --policy <name>    Set configuration policy (stable, compat, preview|nightly). Default: compat
 
 Version specification formats (follows pip format):
     package==2.1.0     Exact version
@@ -66,19 +69,26 @@ def main():
 
     if command == "install":
         args = sys.argv[2:] if len(sys.argv) > 2 else []
-        use_uv = "--uv" in args
-        preview = "--preview" in args
-        unsupported = "--no-unsupported" not in args
-        # Remove flags from args to get package list
-        package_versions = [arg for arg in args if arg not in ("--uv", "--preview", "--no-unsupported")] if args else None
+        try:
+            preview, unsupported, cleaned_args = parse_policy_args(args)
+        except ValueError as e:
+            print(f"Error: {e}")
+            return
+
+        use_uv = "--uv" in cleaned_args
+        # Remove --uv from package list
+        package_versions = [arg for arg in cleaned_args if arg != "--uv"]
         install(package_versions, use_uv=use_uv, preview=preview, unsupported=unsupported)
     elif command == "test":
         subcommand = sys.argv[2] if len(sys.argv) > 2 else "all"
         test(subcommand)
     elif command == "info":
         args = sys.argv[2:] if len(sys.argv) > 2 else []
-        preview = "--preview" in args
-        unsupported = "--no-unsupported" not in args
+        try:
+            preview, unsupported, _ = parse_policy_args(args)
+        except ValueError as e:
+            print(f"Error: {e}")
+            return
         from .utils import info
         info(preview=preview, unsupported=unsupported)
     else:

torchruntime/consts.py

@@ -3,3 +3,10 @@
 AMD = "1002"
 NVIDIA = "10de"
 INTEL = "8086"
+
+POLICIES = {
+    "stable": (False, False),  # preview=False, unsupported=False
+    "compat": (False, True),   # preview=False, unsupported=True (Default)
+    "preview": (True, True),   # preview=True, unsupported=True
+    "nightly": (True, True),   # alias for preview
+}

torchruntime/utils/args.py

@@ -0,0 +1,68 @@
+from ..consts import POLICIES
+
+
+def parse_policy_args(args):
+    """
+    Parses arguments for policy and flags.
+    Returns (preview, unsupported, cleaned_args)
+
+    Supports both `--policy NAME` and `--policy=NAME`.
+
+    Logic:
+    1. Determine base configuration from the LAST provided --policy argument (or default 'compat').
+    2. Apply explicit flags (--preview, --no-unsupported) which ALWAYS override the policy.
+    3. Remove policy and flags from args to produce cleaned_args.
+    """
+    # Default: compat
+    preview, unsupported = POLICIES["compat"]
+    
+    # 1. Scan for the last policy to set the baseline
+    last_policy_name = None
+    i = 0
+    while i < len(args):
+        arg = args[i]
+        if arg == "--policy":
+            if i + 1 < len(args):
+                last_policy_name = args[i+1]
+                i += 2
+            else:
+                # We will catch this error in the second pass or we can raise it now.
+                # Raising now is safer.
+                raise ValueError("--policy requires an argument")
+        elif arg.startswith("--policy="):
+            last_policy_name = arg.split("=", 1)[1]
+            if not last_policy_name:
+                raise ValueError("--policy requires an argument")
+            i += 1
+        else:
+            i += 1
+
+    if last_policy_name:
+        if last_policy_name in POLICIES:
+            preview, unsupported = POLICIES[last_policy_name]
+        else:
+            raise ValueError(f"Unknown policy: {last_policy_name}")
+
+    # 2. Apply flags and build cleaned_args
+    cleaned_args = []
+    i = 0
+    while i < len(args):
+        arg = args[i]
+        if arg == "--policy":
+            # Skip policy and its value (already processed)
+            i += 2
+            continue
+        elif arg.startswith("--policy="):
+            i += 1
+            continue
+        elif arg == "--preview":
+            preview = True
+            i += 1
+        elif arg == "--no-unsupported":
+            unsupported = False
+            i += 1
+        else:
+            cleaned_args.append(arg)
+            i += 1
+            
+    return preview, unsupported, cleaned_args