config.py

import os
import copy

from jetson_containers import L4T_VERSION, CUDA_ARCHITECTURES, github_latest_commit, log_debug

repo = 'mlc-ai/mlc-llm'

package['build_args'] = {
    'MLC_REPO': repo,
    'CUDAARCHS': ';'.join([str(x) for x in CUDA_ARCHITECTURES]),
    'TORCH_CUDA_ARCH_LIST': ';'.join([f'{x/10:.1f}' for x in CUDA_ARCHITECTURES])
}

def mlc(version, patch, llvm=17, tag=None, requires=None, default=False):
    pkg = copy.deepcopy(package)
    
    if default:
        pkg['alias'] = 'mlc'
        
    if requires:
        pkg['requires'] = requires
        
    if not tag:
        tag = version
        
    pkg['name'] = f'mlc:{tag}'

    pkg['build_args'].update({
        'MLC_VERSION': version,
        'MLC_PATCH': patch,
        'LLVM_VERSION': llvm,
    })
    
    pkg['notes'] = f"[{repo}](https://github.com/{repo}/tree/{version}) commit SHA [`{version}`](https://github.com/{repo}/tree/{version})"
    
    return pkg

latest_sha = github_latest_commit(repo, branch='main')
log_debug('-- MLC latest commit:', latest_sha)

#default_dev=(L4T_VERSION.major >= 36)

package = [
    mlc(latest_sha, 'patches/51fb0f4.diff', llvm=18, tag='dev'), #, default=default_dev),
    mlc('9bf5723', 'patches/9bf5723.diff', llvm=17, requires='==35.*'), # 10/20/2023
    mlc('51fb0f4', 'patches/51fb0f4.diff', llvm=17, default=True), # 12/15/2023
]