code/flash-linear-attention/scripts/build_packages.py

#!/usr/bin/env python3
"""Build split packages with proper dependency management and copy them to target directory."""

import ast
import re
import shutil
import subprocess
import sys
from pathlib import Path


def extract_dependencies():
    """Extract dependencies from setup.py using AST."""
    # Get script directory and find setup.py in parent directory
    script_dir = Path(__file__).parent
    setup_py = script_dir.parent / 'setup.py'

    with open(setup_py, encoding='utf-8') as f:
        tree = ast.parse(f.read(), filename=str(setup_py))

    all_deps = []
    extras = {}

    for node in ast.walk(tree):
        if isinstance(node, ast.Call) and getattr(node.func, 'id', '') == 'setup':
            for keyword in node.keywords:
                if (keyword.arg == 'install_requires' and
                        isinstance(keyword.value, (ast.List, ast.Tuple))):
                    all_deps.extend([
                        elt.value for elt in keyword.value.elts
                        if isinstance(elt, ast.Constant) and isinstance(elt.value, str)
                    ])
                elif (keyword.arg == 'extras_require' and
                      isinstance(keyword.value, ast.Dict)):
                    for key_node, val_node in zip(keyword.value.keys, keyword.value.values, strict=False):
                        if (isinstance(key_node, ast.Constant) and
                            isinstance(key_node.value, str) and
                                isinstance(val_node, (ast.List, ast.Tuple))):
                            key = key_node.value
                            values = [
                                elt.value for elt in val_node.elts
                                if isinstance(elt, ast.Constant) and isinstance(elt.value, str)
                            ]
                            extras[key] = values
            break  # Assume only one setup() call

    return all_deps, extras


def categorize_dependencies(deps):
    """Categorize dependencies based on core vs extension."""
    core_deps = []
    ext_deps = []

    for dep in deps:
        if any(core in dep for core in ['torch', 'einops']):
            core_deps.append(dep)
        else:
            ext_deps.append(dep)

    return core_deps, ext_deps


def create_pyproject_toml(package_dir, name, version, dependencies, extras=None):
    """Create pyproject.toml for a package."""
    if extras is None:
        extras = {}

    extras_content = ""
    if extras:
        extras_content = "\n[project.optional-dependencies]\n"
        for key, values in extras.items():
            values_str = ', '.join(f'"{v}"' for v in values)
            extras_content += f"{key} = [{values_str}]\n"

    deps_content = ', '.join(f'"{dep}"' for dep in dependencies)

    # Create description text
    if name == 'fla-core':
        desc_text = 'Core operations for flash-linear-attention'
    else:
        desc_text = 'Fast linear attention models and layers'

    content = f"""[build-system]
requires = ["setuptools", "wheel"]
build-backend = "setuptools.build_meta"

[project]
name = "{name}"
version = "{version}"
description = "{desc_text}"
readme = "README.md"
requires-python = ">=3.10"
dependencies = [{deps_content}]

[project.urls]
Homepage = "https://github.com/fla-org/flash-linear-attention"
Repository = "https://github.com/fla-org/flash-linear-attention"
"""

    content += extras_content

    # Add setuptools namespace package configuration for extension package
    if name == 'flash-linear-attention':
        content += """

[tool.setuptools.packages.find]
include = ["fla*"]
namespaces = true
"""

    with open(package_dir / 'pyproject.toml', 'w') as f:
        f.write(content)


def build_split_packages():
    """Build split packages with proper dependency management."""
    # Get script directory and find files relative to it
    script_dir = Path(__file__).parent
    root_dir = script_dir.parent

    # Get current version
    init_file = root_dir / 'fla' / '__init__.py'
    with open(init_file, encoding='utf-8') as f:
        content = f.read()
    version_match = re.search(r"^__version__\s*=\s*['\"]([^'\"]+)['\"]\s*$", content, re.MULTILINE)
    if not version_match:
        raise RuntimeError(f"Could not find __version__ in {init_file}")
    version = version_match.group(1)

    # Extract dependencies
    all_deps, extras = extract_dependencies()
    core_deps, ext_deps = categorize_dependencies(all_deps)

    # Add version constraint for fla-core in extension package
    ext_deps.insert(0, f'fla-core=={version}')

    # Create output directory
    output_dir = script_dir / 'dist'
    output_dir.mkdir(exist_ok=True)

    # Create fla-core package
    core_dir = output_dir / 'fla-core'
    if core_dir.exists():
        shutil.rmtree(core_dir)
    core_dir.mkdir()

    # Copy core files
    fla_core = core_dir / 'fla'
    shutil.copytree(root_dir / 'fla' / 'ops', fla_core / 'ops')
    shutil.copytree(root_dir / 'fla' / 'modules', fla_core / 'modules')
    shutil.copy(root_dir / 'fla' / 'utils.py', fla_core / 'utils.py')

    # Create fla-core __init__.py
    with open(fla_core / '__init__.py', 'w') as f:
        f.write(f"""# -*- coding: utf-8 -*-

__path__ = __import__('pkgutil').extend_path(__path__, __name__)
__version__ = '{version}'
""")

    # Copy ancillary files (README.md, LICENSE) to core package
    for fname in ("README.md", "LICENSE"):
        src = root_dir / fname
        if src.exists():
            shutil.copy(src, core_dir / fname)

    # Create fla-core configs
    create_pyproject_toml(core_dir, 'fla-core', version, core_deps)

    # Create flash-linear-attention package
    ext_dir = output_dir / 'flash-linear-attention'
    if ext_dir.exists():
        shutil.rmtree(ext_dir)
    ext_dir.mkdir()

    # Copy extension files
    fla_ext = ext_dir / 'fla'
    shutil.copytree(root_dir / 'fla' / 'models', fla_ext / 'models')
    shutil.copytree(root_dir / 'fla' / 'layers', fla_ext / 'layers')

    # Intentionally do NOT create fla/__init__.py in the extension package.
    # The top-level package is provided by fla-core (namespace via pkgutil).

    # Copy ancillary files (README.md, LICENSE) to extension package
    for fname in ("README.md", "LICENSE"):
        src = root_dir / fname
        if src.exists():
            shutil.copy(src, ext_dir / fname)

    # Create extension configs
    create_pyproject_toml(ext_dir, 'flash-linear-attention', version, ext_deps, extras)

    # Create build script
    build_script = output_dir / 'build.sh'
    with open(build_script, 'w') as f:
        f.write("""#!/bin/bash
# Build both packages

echo "Building fla-core..."
cd fla-core
pip install -U build
python -m build

echo "Building flash-linear-attention..."
cd ../flash-linear-attention
python -m build

echo "Build complete! Packages in dist/"
""")

    build_script.chmod(0o755)

    print(f"✅ Split packages created in {output_dir}")
    print(f"✅ fla-core dependencies: {len(core_deps)} packages")
    print(f"✅ flash-linear-attention dependencies: {len(ext_deps)} packages")
    print(f"✅ Version: {version}")

    return output_dir, version


def build_packages(dist_dir):
    """Build wheels and source distributions for both packages."""
    print("Building packages...")

    # Build fla-core (both wheel and sdist)
    print("Building fla-core packages...")
    try:
        subprocess.run(
            [sys.executable, "-m", "build", str(dist_dir / "fla-core")],
            check=True,
            timeout=1800,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
        )
    except subprocess.CalledProcessError as e:
        print("Failed to build fla-core packages:")
        print(e.stdout)
        return False
    except subprocess.TimeoutExpired:
        print("Timed out building fla-core packages")
        return False

    # Build flash-linear-attention (both wheel and sdist)
    print("Building flash-linear-attention packages...")
    try:
        subprocess.run(
            [sys.executable, "-m", "build", str(dist_dir / "flash-linear-attention")],
            check=True,
            timeout=1800,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
        )
    except subprocess.CalledProcessError as e:
        print("Failed to build flash-linear-attention packages:")
        print(e.stdout)
        return False
    except subprocess.TimeoutExpired:
        print("Timed out building flash-linear-attention packages")
        return False

    print("✅ Packages built successfully")
    return True


def copy_packages_to_output(dist_dir):
    """Copy wheels and source distributions to output directory."""
    # Get script directory (relative to this file)
    script_dir = Path(__file__).parent
    root_dir = script_dir.parent

    # Create output directory (relative to root)
    output_dir = root_dir / 'dist-packages'
    output_dir.mkdir(exist_ok=True)

    # Find wheels and source distributions
    core_wheels = list((dist_dir / 'fla-core' / 'dist').glob('*.whl'))
    core_sdist = list((dist_dir / 'fla-core' / 'dist').glob('*.tar.gz'))
    ext_wheels = list((dist_dir / 'flash-linear-attention' / 'dist').glob('*.whl'))
    ext_sdist = list((dist_dir / 'flash-linear-attention' / 'dist').glob('*.tar.gz'))

    if not core_wheels:
        print("No fla-core wheel found")
        return False
    if not ext_wheels:
        print("No flash-linear-attention wheel found")
        return False

    # Copy all packages to output directory
    all_packages = core_wheels + core_sdist + ext_wheels + ext_sdist
    for package in all_packages:
        target = output_dir / package.name
        shutil.copy2(package, target)
        if package.suffix == ".whl":
            package_type = "wheel"
        elif package.suffixes[-2:] == [".tar", ".gz"]:
            package_type = "sdist"
        else:
            package_type = "source"
        print(f"📦 Copied {package_type} package {package.name} to {output_dir}")

    print(f"\n✅ All packages copied to: {output_dir}")
    print("You can install wheels with:")
    print("  pip install dist-packages/*.whl")
    print("Source distributions are also available in:", output_dir)

    return True


def main():
    """Build split packages and copy to target directory."""

    print("Building split packages...")

    # Build the split packages
    dist_dir, _ = build_split_packages()

    print("\nTo build packages manually:")
    print(f"cd {dist_dir}")
    print("./build.sh")

    # Build packages (wheels and source distributions)
    if not build_packages(dist_dir):
        return 1

    # Copy packages to output directory
    if not copy_packages_to_output(dist_dir):
        return 1

    return 0


if __name__ == "__main__":
    exit(main())