File size: 1,633 Bytes
750aacd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 | """Exercise launcher arguments without loading models or device SDKs."""
from pathlib import Path
import os,subprocess,tempfile,unittest
ROOT=Path(__file__).resolve().parents[1]
class LauncherTest(unittest.TestCase):
def run_launcher(self,backend,override=None):
with tempfile.TemporaryDirectory() as tmp:
fake=Path(tmp)/'python'
fake.write_text('#!/usr/bin/env bash\nif [[ "$1" == "-" ]]; then exit 0; fi\nprintf "%s\\n" "$@"\nprintf "TOKENIZER_DIR=%s\\n" "${TOKENIZER_DIR:-}"\n')
fake.chmod(0o755)
env={'PATH':os.environ['PATH'],'UMI_INFER_DATA':'/cell','FLUX_RESOURCES':'/encoders','FLUX_PYTHON':str(fake),'PI05_PYTHON':str(fake)}
if override:env.update(override)
command=['bash','-c','TOKENIZER_DIR=/tokenizer; launcher="$1"; shift; source "$launcher"','test',str(ROOT/f'scripts/run_{backend}_server.sh')]
return subprocess.check_output(command,env=env,text=True).splitlines()
def test_flux_defaults_and_override(self):
args=self.run_launcher('flux');i=args.index('--flux-export')
self.assertEqual(args[i+1],str(ROOT/'checkpoints/flux_fold_clothes_bf16'))
args=self.run_launcher('flux',{'FLUX_EXPORT_DIR':'/other/export'})
self.assertEqual(args[args.index('--flux-export')+1],'/other/export')
def test_pi05_default_and_tokenizer_export(self):
args=self.run_launcher('pi05')
self.assertEqual(args[args.index('--checkpoint-dir')+1],str(ROOT/'checkpoints/pi05_fold_clothes/epoch_2'))
self.assertIn('TOKENIZER_DIR=/tokenizer',args)
if __name__=='__main__':unittest.main()
|