icarus112 commited on
Commit
6ef2d06
·
verified ·
1 Parent(s): 313e3b0

feat(triton-cache): wire setup/teardown into entrypoint

Browse files
Files changed (1) hide show
  1. 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: