Spaces:
Runtime error
Runtime error
feat(triton-cache): wire setup/teardown into entrypoint
Browse files- entrypoint.py +15 -0
entrypoint.py
CHANGED
|
@@ -60,6 +60,14 @@ def _early_cuda_kick() -> None:
|
|
| 60 |
|
| 61 |
_early_cuda_kick()
|
| 62 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 63 |
from huggingface_hub import HfApi # noqa: E402 (import after cuda kick)
|
| 64 |
|
| 65 |
REPO_ROOT = Path('/workspace/feather')
|
|
@@ -172,6 +180,13 @@ def run_job_mode() -> int:
|
|
| 172 |
print(f'[job] command={cmd}', flush=True)
|
| 173 |
proc = subprocess.run(cmd, check=False)
|
| 174 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 175 |
if TOKEN:
|
| 176 |
api = HfApi(token=TOKEN)
|
| 177 |
try:
|
|
|
|
| 60 |
|
| 61 |
_early_cuda_kick()
|
| 62 |
|
| 63 |
+
# Hydrate triton compilation cache from HF Hub before any triton/mamba_ssm import.
|
| 64 |
+
# triton_cache_setup.py is copied next to this file by the job bash command.
|
| 65 |
+
try:
|
| 66 |
+
import triton_cache_setup as _tcs
|
| 67 |
+
_tcs.setup()
|
| 68 |
+
except ImportError:
|
| 69 |
+
print('[boot] triton_cache_setup not found; skipping cache hydrate', flush=True)
|
| 70 |
+
|
| 71 |
from huggingface_hub import HfApi # noqa: E402 (import after cuda kick)
|
| 72 |
|
| 73 |
REPO_ROOT = Path('/workspace/feather')
|
|
|
|
| 180 |
print(f'[job] command={cmd}', flush=True)
|
| 181 |
proc = subprocess.run(cmd, check=False)
|
| 182 |
|
| 183 |
+
# Push triton compilation cache back to HF Hub for next run.
|
| 184 |
+
try:
|
| 185 |
+
import triton_cache_setup as _tcs
|
| 186 |
+
_tcs.teardown()
|
| 187 |
+
except Exception as _tcs_err:
|
| 188 |
+
print(f'[triton_cache] teardown error (non-fatal): {_tcs_err}', flush=True)
|
| 189 |
+
|
| 190 |
if TOKEN:
|
| 191 |
api = HfApi(token=TOKEN)
|
| 192 |
try:
|