1818        "PATH" : "/root/.cargo/bin:$PATH" ,
1919        "BUILDKITE_REPO" : os .environ .get ("BUILDKITE_REPO" , "" ),
2020        "BUILDKITE_COMMIT" : os .environ .get ("BUILDKITE_COMMIT" , "" ),
21+         "BUILDKITE_PULL_REQUEST" : os .environ .get ("BUILDKITE_PULL_REQUEST" , "" ),
22+         "IMAGE_VERSION" : os .environ .get ("IMAGE_VERSION" , "" ),
2123    })
2224)
2325
@@ -29,16 +31,27 @@ def run_test(pytest_command: str):
2931
3032    git_repo  =  os .environ .get ("BUILDKITE_REPO" )
3133    git_commit  =  os .environ .get ("BUILDKITE_COMMIT" )
34+     pr_number  =  os .environ .get ("BUILDKITE_PULL_REQUEST" )
3235
3336    print (f"Cloning repository: { git_repo }  " )
34-     print (f"Checking out commit: { git_commit }  " )
37+     print (f"Target commit: { git_commit }  " )
38+     if  pr_number :
39+         print (f"PR number: { pr_number }  " )
40+     
41+     # For PRs (including forks), use GitHub's PR refs to get the correct commit 
42+     if  pr_number  and  pr_number  !=  "false" :
43+         checkout_command  =  f"git fetch --prune origin refs/pull/{ pr_number }  /head && git checkout FETCH_HEAD" 
44+         print (f"Using PR ref for checkout: { checkout_command }  " )
45+     else :
46+         checkout_command  =  f"git checkout { git_commit }  " 
47+         print (f"Using direct commit checkout: { checkout_command }  " )
3548
3649    command  =  f""" 
3750    source $HOME/.local/bin/env && 
3851    source /opt/venv/bin/activate && 
3952    git clone { git_repo }   /FastVideo && 
4053    cd /FastVideo && 
41-     git checkout  { git_commit }   && 
54+     { checkout_command }   && 
4255    uv pip install -e .[test] && 
4356    { pytest_command }  
4457    """ 
@@ -49,38 +62,38 @@ def run_test(pytest_command: str):
4962
5063    sys .exit (result .returncode )
5164
52- @app .function (gpu = "L40S:1" , image = image , timeout = 1800 ) 
65+ @app .function (gpu = "L40S:1" , image = image , timeout = 900 ) 
5366def  run_encoder_tests ():
5467    run_test ("pytest ./fastvideo/v1/tests/encoders -vs" )
5568
56- @app .function (gpu = "L40S:1" , image = image , timeout = 1800 ) 
69+ @app .function (gpu = "L40S:1" , image = image , timeout = 900 ) 
5770def  run_vae_tests ():
5871    run_test ("pytest ./fastvideo/v1/tests/vaes -vs" )
5972
60- @app .function (gpu = "L40S:1" , image = image , timeout = 1800 ) 
73+ @app .function (gpu = "L40S:1" , image = image , timeout = 900 ) 
6174def  run_transformer_tests ():
6275    run_test ("pytest ./fastvideo/v1/tests/transformers -vs" )
6376
64- @app .function (gpu = "L40S:2" , image = image , timeout = 3600 ) 
77+ @app .function (gpu = "L40S:2" , image = image , timeout = 1800 ) 
6578def  run_ssim_tests ():
6679    run_test ("pytest ./fastvideo/v1/tests/ssim -vs" )
6780
68- @app .function (gpu = "L40S:4" , image = image , timeout = 1800 , secrets = [modal .Secret .from_dict ({"WANDB_API_KEY" : os .environ .get ("WANDB_API_KEY" , "" )})]) 
81+ @app .function (gpu = "L40S:4" , image = image , timeout = 900 , secrets = [modal .Secret .from_dict ({"WANDB_API_KEY" : os .environ .get ("WANDB_API_KEY" , "" )})]) 
6982def  run_training_tests ():
7083    run_test ("wandb login $WANDB_API_KEY && pytest ./fastvideo/v1/tests/training/Vanilla -srP" )
7184
72- @app .function (gpu = "H100:2" , image = image , timeout = 1800 , secrets = [modal .Secret .from_dict ({"WANDB_API_KEY" : os .environ .get ("WANDB_API_KEY" , "" )})]) 
85+ @app .function (gpu = "H100:2" , image = image , timeout = 900 , secrets = [modal .Secret .from_dict ({"WANDB_API_KEY" : os .environ .get ("WANDB_API_KEY" , "" )})]) 
7386def  run_training_tests_VSA ():
7487    run_test ("wandb login $WANDB_API_KEY && pytest ./fastvideo/v1/tests/training/VSA -srP" )
7588
76- @app .function (gpu = "H100:2" , image = image , timeout = 1800 ) 
89+ @app .function (gpu = "H100:2" , image = image , timeout = 900 ) 
7790def  run_inference_tests_STA ():
7891    run_test ("pytest ./fastvideo/v1/tests/inference/STA -srP" )
7992
80- @app .function (gpu = "H100:1" , image = image , timeout = 1800 ) 
93+ @app .function (gpu = "H100:1" , image = image , timeout = 900 ) 
8194def  run_precision_tests_STA ():
8295    run_test ("python csrc/attn/tests/test_sta.py" )
8396
84- @app .function (gpu = "H100:1" , image = image , timeout = 1800 ) 
97+ @app .function (gpu = "H100:1" , image = image , timeout = 900 ) 
8598def  run_precision_tests_VSA ():
8699    run_test ("python csrc/attn/tests/test_block_sparse.py" )
0 commit comments