systemone.py

"""Run one System One request, or a batch of requests, on a local Decision model."""
import argparse
import json
from pathlib import Path

NATIVE_MANIFEST_SHA256 = 'f288d873999832a3f37c6a7c4268c2ab309691e621794dbf7acab891acbbb7e6'


def _unique_object(pairs):
    result = {}
    for key, value in pairs:
        if key in result:
            raise ValueError('Duplicate JSON object key: ' + key)
        result[key] = value
    return result


def main():
    parser = argparse.ArgumentParser(description='Decision System One: typed parallel questions on AMD')
    parser.add_argument('--native', default=str(Path(__file__).resolve().parent / 'native'))
    parser.add_argument('--manifest-sha256', default=NATIVE_MANIFEST_SHA256)
    parser.add_argument('--model', help='Explicit name for a compatible custom fine-tune')
    parser.add_argument('--input', required=True, help='JSON request or array of request objects')
    parser.add_argument('--output', required=True, help='Fresh response JSON file')
    parser.add_argument('--batching', choices=('default', 'auto'), default='default')
    args = parser.parse_args()
    payload = Path(args.input).read_bytes()
    if len(payload) > 2 * 1024 * 1024:
        raise ValueError('Input file exceeds 2 MiB')
    request = json.loads(payload, object_pairs_hook=_unique_object)
    if Path(args.output).exists():
        raise ValueError('Output must be fresh')
    import torch
    from decision_runtime import load_native
    from decision_inference import SystemOne
    if torch.version.hip is None or not torch.cuda.is_available() or torch.cuda.device_count() != 1:
        raise RuntimeError('Expose exactly one AMD ROCm GPU')
    torch.cuda.set_device(0)
    torch.set_num_threads(2)
    torch.backends.cuda.matmul.allow_tf32 = False
    torch.backends.cudnn.allow_tf32 = False
    torch.backends.mha.set_fastpath_enabled(False)
    native = load_native(args.native, expected_manifest_sha256=args.manifest_sha256, device='cuda:0')
    client = SystemOne(native, model=args.model, batching=args.batching)
    result = client.batch(request) if isinstance(request, list) else client.evaluate(request)
    with Path(args.output).open('x', encoding='utf-8') as stream:
        json.dump(result, stream, ensure_ascii=False, allow_nan=False, indent=2)
        stream.write('\n')


if __name__ == '__main__':
    main()