tests/test_gpu_smoke.py

import os
import unittest


class GPUCheckpointSmokeTest(unittest.TestCase):
    @unittest.skipUnless(
        os.getenv("PI05_GPU_SMOKE") == "1",
        "set PI05_GPU_SMOKE=1 to run the real checkpoint test",
    )
    def test_real_checkpoint_prediction(self):
        import numpy as np
        from PIL import Image

        from artifacts import resolve_checkpoint_path, resolve_model_id
        from inference import run_prediction
        from model_loader import DEFAULT_POLICY_CONFIG, MODEL_MANAGER

        model_id = resolve_model_id()
        checkpoint = resolve_checkpoint_path()
        policy = MODEL_MANAGER.get(model_id, checkpoint, DEFAULT_POLICY_CONFIG)
        image = Image.fromarray(np.zeros((224, 224, 3), dtype=np.uint8))
        result = run_prediction(
            policy,
            image,
            image,
            "move safely",
            [0.0] * 7,
            0,
            model_id,
            checkpoint,
        )
        self.assertEqual(result.actions.shape, (10, 7))


if __name__ == "__main__":
    unittest.main()