File size: 924 Bytes
d9bb75c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 | # Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This software may be used and distributed in accordance with
# the terms of the DINOv3 License Agreement.
import contextlib
import logging
import os
import tempfile
from typing import Optional
import submitit.helpers
logger = logging.getLogger("dinov3")
@contextlib.contextmanager
def clean_env():
try:
# Hide torch.compile() variables from the launched evals
extra_names = ("TRITON_CACHE_DIR", "TORCHINDUCTOR_CACHE_DIR")
ctx = submitit.helpers.clean_env(extra_names=extra_names)
except TypeError as e:
logger.warning("Update submitit to the latest main branch\n%s", e)
ctx = submitit.helpers.clean_env()
with ctx:
yield
def set_triton_cache_dir(cache_dir: Optional[str] = None) -> None:
if cache_dir is None:
cache_dir = tempfile.mkdtemp()
os.environ["TRITON_CACHE_DIR"] = cache_dir
|