diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 9c266a26fc..048aeeafdc 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -77,6 +77,38 @@ def _model_wants_causal_conv1d(model_name: str) -> bool: ) +def _hipcc_gcc_install_dir() -> str | None: + """Return the highest-numbered ``/usr/lib/gcc/x86_64-linux-gnu/`` that has + BOTH the gcc runtime dir AND the corresponding ``/usr/include/c++/`` C++ + headers, or ``None`` if no match (or non-Linux / non-x86_64). + + Ubuntu 24.04 ships ``/usr/lib/gcc/x86_64-linux-gnu/14/`` (gcc-14 runtime + objects) but does NOT ship ``/usr/include/c++/14`` in its default apt set; + libstdc++ headers come from ``libstdc++-13-dev``. ROCm clang-20 picks the + highest-numbered runtime dir by default, finds no ````, and the + HIP source build fails with:: + + /opt/rocm-X.Y/lib/llvm/lib/clang/20/include/__clang_hip_runtime_wrapper.h:112:10: + fatal error: 'cstdlib' file not found + + Returning a path lets the caller pass ``--gcc-install-dir=`` to clang + via ``HIPCC_COMPILE_FLAGS_APPEND``. Mirrors the same loop ``bbf004c`` added + to ``studio/setup.sh`` for the llama.cpp HIP build branch (PR #5301). + """ + if not sys.platform.startswith("linux"): + return None + import platform as _platform + + if _platform.machine().lower() != "x86_64": + return None + for _ver in (14, 13, 12, 11): + _runtime = f"/usr/lib/gcc/x86_64-linux-gnu/{_ver}/include" + _headers = f"/usr/include/c++/{_ver}" + if os.path.isdir(_runtime) and os.path.isdir(_headers): + return f"/usr/lib/gcc/x86_64-linux-gnu/{_ver}" + return None + + def _install_package_wheel_first( *, event_queue: Any, @@ -212,6 +244,30 @@ def _install_package_wheel_first( } if is_hip: _run_kwargs["timeout"] = 1800 + # On Ubuntu 24.04 + ROCm clang-20, the HIP source build (causal-conv1d, + # mamba-ssm source fallback, flash-attn source fallback) defaults to + # /usr/lib/gcc/x86_64-linux-gnu/14/ which has the runtime dir but no + # /usr/include/c++/14 headers, and dies at: + # __clang_hip_runtime_wrapper.h:112:10: + # fatal error: 'cstdlib' file not found + # Inject --gcc-install-dir for a gcc whose C++ headers actually exist. + # Respect any pre-existing --gcc-install-dir in HIPCC_COMPILE_FLAGS_APPEND + # (user knows best); otherwise append. Mirrors the same fix bbf004c + # added to studio/setup.sh for the llama.cpp HIP build (PR #5301). + _existing_flags = os.environ.get("HIPCC_COMPILE_FLAGS_APPEND", "") + if "--gcc-install-dir" not in _existing_flags: + _gcc_dir = _hipcc_gcc_install_dir() + if _gcc_dir is not None: + _appended = (f"{_existing_flags} --gcc-install-dir={_gcc_dir}").strip() + _env = _run_kwargs.get("env", os.environ).copy() + _env["HIPCC_COMPILE_FLAGS_APPEND"] = _appended + _run_kwargs["env"] = _env + logger.info( + "HIP source build for %s: appended " + "--gcc-install-dir=%s to HIPCC_COMPILE_FLAGS_APPEND", + display_name, + _gcc_dir, + ) try: result = _sp.run(pypi_cmd, **_run_kwargs) diff --git a/studio/backend/tests/test_training_worker_flash_attn.py b/studio/backend/tests/test_training_worker_flash_attn.py index 0737bdc82f..733b726656 100644 --- a/studio/backend/tests/test_training_worker_flash_attn.py +++ b/studio/backend/tests/test_training_worker_flash_attn.py @@ -6,6 +6,7 @@ from __future__ import annotations import builtins import subprocess import sys +from typing import Any from unittest import mock from core.training import worker @@ -22,6 +23,17 @@ def _missing_flash_attn_import(): return fake_import +def _missing_module_import(missing: str): + real_import = builtins.__import__ + + def fake_import(name, globals = None, locals = None, fromlist = (), level = 0): + if name == missing: + raise ImportError + return real_import(name, globals, locals, fromlist, level) + + return fake_import + + def test_should_try_runtime_flash_attn_install_threshold_and_skip(monkeypatch): monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False) assert worker._should_try_runtime_flash_attn_install(32767) is False @@ -193,3 +205,270 @@ def test_mamba_ssm_path_preserves_wheel_first_install_args(monkeypatch): release_tag = worker._MAMBA_SSM_RELEASE_TAG, release_base_url = "https://github.com/state-spaces/mamba/releases/download", ) + + +# ──────────────────────────────────────────────────────────────────── +# HIP source-build gcc-install-dir coverage (h34v3nzc0dex Strix Halo). +# Ubuntu 24.04 ships gcc-14's runtime dir without /usr/include/c++/14, +# so ROCm clang-20 picks it and fails with 'cstdlib' file not found +# when building causal-conv1d (or any other HIP source fallback). +# _hipcc_gcc_install_dir() finds a gcc dir that has both halves; the +# _install_package_wheel_first HIP branch passes it to clang via +# HIPCC_COMPILE_FLAGS_APPEND. Parallel to bbf004c's setup.sh fix for +# the llama.cpp HIP build (PR #5301). +# ──────────────────────────────────────────────────────────────────── + + +def _isdir_for_layout(*existing: str): + """Return an os.path.isdir replacement that only treats the given + absolute paths as directories. Lets a test simulate exactly which + gcc runtime dirs and C++ header dirs exist on the host.""" + valid = set(existing) + + def fake_isdir(path: str) -> bool: + return path in valid + + return fake_isdir + + +def test_hipcc_gcc_install_dir_picks_highest_with_headers(monkeypatch): + """gcc-14 has runtime but no /usr/include/c++/14; loop falls through + to gcc-13 which has both. This is the exact Ubuntu 24.04 layout.""" + monkeypatch.setattr(sys, "platform", "linux") + import platform as _platform + + monkeypatch.setattr(_platform, "machine", lambda: "x86_64") + monkeypatch.setattr( + worker.os.path, + "isdir", + _isdir_for_layout( + "/usr/lib/gcc/x86_64-linux-gnu/14/include", # runtime present + # but no /usr/include/c++/14 — typical Ubuntu 24.04 default + "/usr/lib/gcc/x86_64-linux-gnu/13/include", + "/usr/include/c++/13", # libstdc++-13-dev installed + ), + ) + assert worker._hipcc_gcc_install_dir() == "/usr/lib/gcc/x86_64-linux-gnu/13" + + +def test_hipcc_gcc_install_dir_picks_14_when_headers_exist(monkeypatch): + """If the user has libstdc++-14-dev installed, prefer gcc-14.""" + monkeypatch.setattr(sys, "platform", "linux") + import platform as _platform + + monkeypatch.setattr(_platform, "machine", lambda: "x86_64") + monkeypatch.setattr( + worker.os.path, + "isdir", + _isdir_for_layout( + "/usr/lib/gcc/x86_64-linux-gnu/14/include", + "/usr/include/c++/14", + ), + ) + assert worker._hipcc_gcc_install_dir() == "/usr/lib/gcc/x86_64-linux-gnu/14" + + +def test_hipcc_gcc_install_dir_returns_none_when_no_match(monkeypatch): + """No gcc dir has both halves → return None and skip the env injection + rather than guessing wrong and surfacing a confusing build failure.""" + monkeypatch.setattr(sys, "platform", "linux") + import platform as _platform + + monkeypatch.setattr(_platform, "machine", lambda: "x86_64") + monkeypatch.setattr(worker.os.path, "isdir", lambda path: False) + assert worker._hipcc_gcc_install_dir() is None + + +def test_hipcc_gcc_install_dir_returns_none_on_non_linux(monkeypatch): + """Don't probe gcc layout on macOS / Windows — early-return.""" + monkeypatch.setattr(sys, "platform", "darwin") + + def _isdir_should_not_be_called(_path): + raise AssertionError("isdir should not be called on non-Linux") + + monkeypatch.setattr(worker.os.path, "isdir", _isdir_should_not_be_called) + assert worker._hipcc_gcc_install_dir() is None + + +def test_hipcc_gcc_install_dir_returns_none_on_non_x86_64(monkeypatch): + """ROCm clang-20 on aarch64 has a different libstdc++ layout.""" + monkeypatch.setattr(sys, "platform", "linux") + import platform as _platform + + monkeypatch.setattr(_platform, "machine", lambda: "aarch64") + assert worker._hipcc_gcc_install_dir() is None + + +def _make_hip_install_env(monkeypatch, *, gcc_dir: str | None): + """Common scaffolding for tests that exercise the HIP source-build + branch of _install_package_wheel_first end-to-end. The package isn't + installed yet, no prebuilt wheel exists, hipcc is on PATH, and the + fake env reports an HIP torch.""" + monkeypatch.setattr(builtins, "__import__", _missing_module_import("causal_conv1d")) + monkeypatch.setattr( + worker, + "probe_torch_wheel_env", + lambda timeout = 30: { + "hip_version": "7.13.26176", + "python_tag": "cp312", + "torch_mm": "2.11", + "cxx11abi": "TRUE", + "platform_tag": "linux_x86_64", + }, + ) + monkeypatch.setattr(worker, "direct_wheel_url", lambda **kw: None) + monkeypatch.setattr( + worker.shutil, + "which", + lambda name: "/opt/rocm/bin/hipcc" if name == "hipcc" else None, + ) + monkeypatch.setattr(worker, "_send_status", lambda *a, **k: None) + monkeypatch.setattr(worker, "_hipcc_gcc_install_dir", lambda: gcc_dir) + + +def test_install_injects_gcc_install_dir_on_hip_source_build(monkeypatch): + """HIP source-build with no user-set HIPCC_COMPILE_FLAGS_APPEND → + subprocess env carries --gcc-install-dir=.""" + monkeypatch.delenv("HIPCC_COMPILE_FLAGS_APPEND", raising = False) + _make_hip_install_env(monkeypatch, gcc_dir = "/usr/lib/gcc/x86_64-linux-gnu/13") + + captured: dict[str, str] = {} + + def fake_run(cmd, **kwargs): + captured.update(kwargs.get("env") or {}) + return subprocess.CompletedProcess(cmd, 0, "") + + monkeypatch.setattr(worker._sp, "run", fake_run) + + worker._install_package_wheel_first( + event_queue = [], + import_name = "causal_conv1d", + display_name = "causal-conv1d", + pypi_name = "causal-conv1d", + pypi_version = "1.6.2.post1", + filename_prefix = "causal_conv1d", + release_tag = "v1.6.2.post1", + release_base_url = "https://example.com", + ) + + assert ( + captured.get("HIPCC_COMPILE_FLAGS_APPEND") + == "--gcc-install-dir=/usr/lib/gcc/x86_64-linux-gnu/13" + ) + + +def test_install_appends_to_existing_hipcc_compile_flags(monkeypatch): + """User has HIPCC_COMPILE_FLAGS_APPEND='-O3 -DFOO' set → final value + keeps the user's flags AND adds --gcc-install-dir at the end.""" + monkeypatch.setenv("HIPCC_COMPILE_FLAGS_APPEND", "-O3 -DFOO") + _make_hip_install_env(monkeypatch, gcc_dir = "/usr/lib/gcc/x86_64-linux-gnu/13") + + captured: dict[str, str] = {} + + def fake_run(cmd, **kwargs): + captured.update(kwargs.get("env") or {}) + return subprocess.CompletedProcess(cmd, 0, "") + + monkeypatch.setattr(worker._sp, "run", fake_run) + + worker._install_package_wheel_first( + event_queue = [], + import_name = "causal_conv1d", + display_name = "causal-conv1d", + pypi_name = "causal-conv1d", + pypi_version = "1.6.2.post1", + filename_prefix = "causal_conv1d", + release_tag = "v1.6.2.post1", + release_base_url = "https://example.com", + ) + + assert captured.get("HIPCC_COMPILE_FLAGS_APPEND") == ( + "-O3 -DFOO --gcc-install-dir=/usr/lib/gcc/x86_64-linux-gnu/13" + ) + + +def test_install_respects_user_gcc_install_dir(monkeypatch): + """User explicitly set --gcc-install-dir=… already → don't touch it. + Avoids two competing --gcc-install-dir flags on the clang command line.""" + monkeypatch.setenv( + "HIPCC_COMPILE_FLAGS_APPEND", + "--gcc-install-dir=/opt/custom/gcc-13", + ) + _make_hip_install_env(monkeypatch, gcc_dir = "/usr/lib/gcc/x86_64-linux-gnu/13") + + captured: dict[str, str] | None = {"_called": "no"} + + def fake_run(cmd, **kwargs): + env = kwargs.get("env") + if env is not None: + captured.clear() + captured.update(env) + else: + captured["_called"] = "yes_no_env" + return subprocess.CompletedProcess(cmd, 0, "") + + monkeypatch.setattr(worker._sp, "run", fake_run) + + worker._install_package_wheel_first( + event_queue = [], + import_name = "causal_conv1d", + display_name = "causal-conv1d", + pypi_name = "causal-conv1d", + pypi_version = "1.6.2.post1", + filename_prefix = "causal_conv1d", + release_tag = "v1.6.2.post1", + release_base_url = "https://example.com", + ) + + # subprocess.run was invoked without env override (the user already + # set HIPCC_COMPILE_FLAGS_APPEND with --gcc-install-dir, so we left + # the env alone — the existing value is inherited normally). + assert captured == {"_called": "yes_no_env"} + + +def test_install_does_not_inject_env_on_cuda(monkeypatch): + """CUDA path (no hip_version in env) → no env override at all.""" + monkeypatch.delenv("HIPCC_COMPILE_FLAGS_APPEND", raising = False) + monkeypatch.setattr(builtins, "__import__", _missing_module_import("causal_conv1d")) + monkeypatch.setattr( + worker, + "probe_torch_wheel_env", + lambda timeout = 30: { + "python_tag": "cp312", + "torch_mm": "2.11", + "cuda_major": "12", + "cxx11abi": "TRUE", + "platform_tag": "linux_x86_64", + }, + ) + monkeypatch.setattr(worker, "direct_wheel_url", lambda **kw: None) + monkeypatch.setattr(worker.shutil, "which", lambda name: None) + monkeypatch.setattr(worker, "_send_status", lambda *a, **k: None) + # If _hipcc_gcc_install_dir were called on CUDA we'd want to know. + monkeypatch.setattr( + worker, + "_hipcc_gcc_install_dir", + lambda: (_ for _ in ()).throw(AssertionError("must not run on CUDA")), + ) + + captured: dict[str, Any] = {} + + def fake_run(cmd, **kwargs): + captured["env_in_kwargs"] = "env" in kwargs + return subprocess.CompletedProcess(cmd, 0, "") + + monkeypatch.setattr(worker._sp, "run", fake_run) + + worker._install_package_wheel_first( + event_queue = [], + import_name = "causal_conv1d", + display_name = "causal-conv1d", + pypi_name = "causal-conv1d", + pypi_version = "1.6.2.post1", + filename_prefix = "causal_conv1d", + release_tag = "v1.6.2.post1", + release_base_url = "https://example.com", + ) + + # CUDA branch never sets the env, never invokes the gcc helper. + assert captured.get("env_in_kwargs") is False