| """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() |
|
|
