chenlin
f4b415c54f
web-platform-ci / TypeScript、Lint、Unit、Build (push) Has been cancelled
web-platform-ci / Playwright E2E (push) Has been cancelled
web-platform-ci / TypeScript、Lint、Unit、Build (pull_request) Has been cancelled
web-platform-ci / Playwright E2E (pull_request) Has been cancelled
81 lines
2.7 KiB
Python
81 lines
2.7 KiB
Python
import sys
|
|
import tempfile
|
|
import time
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
from server import ApiError, TrainingManager # noqa: E402
|
|
|
|
|
|
class TrainingManagerTest(unittest.TestCase):
|
|
def setUp(self):
|
|
self.temporary = tempfile.TemporaryDirectory()
|
|
self.root = Path(self.temporary.name)
|
|
(self.root / "scripts").mkdir()
|
|
(self.root / "scripts" / "train.py").write_text(
|
|
"""import os, pathlib, time
|
|
print('WANDB_MODE=' + os.environ.get('WANDB_MODE', ''), flush=True)
|
|
print('Learning iteration 1 / 2', flush=True)
|
|
time.sleep(0.02)
|
|
print('Learning iteration 2 / 2', flush=True)
|
|
out=pathlib.Path('logs/rsl_rl/test/run/policy.onnx')
|
|
out.parent.mkdir(parents=True, exist_ok=True)
|
|
out.write_bytes(b'onnx')
|
|
""",
|
|
encoding="utf-8",
|
|
)
|
|
self.manager = TrainingManager(self.root, sys.executable, ("Unitree-Go2-Flat",), check_environment=False)
|
|
|
|
def tearDown(self):
|
|
self.temporary.cleanup()
|
|
|
|
@staticmethod
|
|
def payload(**patch):
|
|
value = {
|
|
"taskId": "Unitree-Go2-Flat",
|
|
"numEnvs": 16,
|
|
"maxIterations": 2,
|
|
"seed": 42,
|
|
"runName": "browser-test",
|
|
"device": "cpu",
|
|
"gpuIds": [],
|
|
"wandbMode": "offline",
|
|
}
|
|
value.update(patch)
|
|
return value
|
|
|
|
def test_validates_allowlist_and_limits(self):
|
|
with self.assertRaises(ApiError):
|
|
self.manager.parse_config(self.payload(taskId="shell injection"))
|
|
with self.assertRaises(ApiError):
|
|
self.manager.parse_config(self.payload(numEnvs=0))
|
|
with self.assertRaises(ApiError):
|
|
self.manager.parse_config(self.payload(runName="bad name"))
|
|
with self.assertRaises(ApiError):
|
|
self.manager.parse_config(self.payload(wandbMode="login"))
|
|
|
|
def test_builds_argument_array_without_shell(self):
|
|
config = self.manager.parse_config(self.payload(device="gpu", gpuIds=[0, 2]))
|
|
command = self.manager.command_for(config)
|
|
self.assertEqual(command[:4], [sys.executable, "-u", "scripts/train.py", "Unitree-Go2-Flat"])
|
|
self.assertEqual(command[-2:], ["--gpu-ids", "[0,2]"])
|
|
|
|
def test_runs_job_and_exposes_new_onnx_artifact(self):
|
|
job = self.manager.start(self.payload())
|
|
deadline = time.monotonic() + 5
|
|
while time.monotonic() < deadline:
|
|
job = self.manager.get(job["id"])
|
|
if job["state"] not in ("queued", "running"):
|
|
break
|
|
time.sleep(0.02)
|
|
self.assertEqual(job["state"], "succeeded")
|
|
self.assertEqual(job["iteration"], 2)
|
|
self.assertIn("WANDB_MODE=offline", job["logs"])
|
|
self.assertTrue(job["artifactReady"])
|
|
self.assertEqual(self.manager.artifact(job["id"]).read_bytes(), b"onnx")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|