from ops import *
ROOT=Path('/data/contextra/model-artifacts/deployments/batch8-b286446')
ROOT.mkdir(mode=0o700,exist_ok=True)
old=json.loads((PRIVATE/'analytics-inspect.json').read_text())
mounts={Path(m['Destination']).name:Path(m['Source']) for m in old['Mounts'] if '/models/triton/' in m['Destination']}
result={}
parser='''import sys
from google.protobuf import text_format
from tritonclient.grpc import model_config_pb2
c=model_config_pb2.ModelConfig();text_format.Parse(sys.stdin.read(),c)
assert c.max_batch_size==4 and list(c.dynamic_batching.preferred_batch_size) in ([],[4])
assert all(g.count==8 for g in c.instance_group)
c.max_batch_size=8;del c.dynamic_batching.preferred_batch_size[:];c.dynamic_batching.preferred_batch_size.append(8)
del c.version_policy.specific.versions[:];c.version_policy.specific.versions.append(8)
print(text_format.MessageToString(c))'''
for model,prior in mounts.items():
    if model=='Awiros_person_and_head_detection':
        candidate=Path('/data/contextra/model-artifacts/candidates/crowd-batch8-b286446')
        expected=json.loads((HERE/'crowd-build-result.json').read_text())['engine_sha256']
    else:
        candidate=Path('/data/contextra/model-artifacts/candidates/batch8-6367102')/model
        manifest=json.loads((candidate/'manifest.json').read_text())
        assert manifest['status']=='prepared_and_execution_checked' and manifest['profile']==[1,8,8]
        expected=manifest['engine_sha256']
    engine=candidate/'8/model.plan';assert hashlib.sha256(engine.read_bytes()).hexdigest()==expected
    folder=ROOT/model;(folder/'8').mkdir(parents=True,exist_ok=True)
    if (folder/'8/model.plan').exists():
        assert hashlib.sha256((folder/'8/model.plan').read_bytes()).hexdigest()==expected
    else:
        shutil.copyfile(engine,folder/'8/model.plan')
    config=cmd(['docker','exec','-i','contextra-analytics-1','python3','-B','-c',parser],input=(prior/'config.pbtxt').read_text())
    (folder/'config.pbtxt').write_text(config)
    result[model]=dict(source=str(candidate),deployment=str(folder),version=8,max_batch_size=8,preferred_batch_size=[8],instances=8,
        engine_sha256=expected,config_sha256=hashlib.sha256(config.encode()).hexdigest(),previous=str(prior))
save('models-prepared.json',result)
print('Four version8/batch8/instance8 deployments prepared')
