raylim commited on
Commit
cbb7db9
·
unverified ·
1 Parent(s): 24ea926

Fix T4 detection by checking actual GPU name

Browse files

- No SPACE_HARDWARE env var on HF Spaces T4
- Detect T4 by checking torch.cuda.get_device_name()
- Look for 'T4' in GPU name (Tesla T4)
- Add hardware logging at module load
- Should now correctly use 4 workers on T4

Files changed (1) hide show
  1. src/mosaic/analysis.py +15 -7
src/mosaic/analysis.py CHANGED
@@ -26,8 +26,16 @@ except ImportError:
26
  return lambda f: f
27
  return fn
28
 
29
- # Detect T4 hardware
30
- IS_T4_GPU = os.environ.get("SPACE_HARDWARE") == "t4-small"
 
 
 
 
 
 
 
 
31
 
32
  # Set optimal parameters based on hardware
33
  if IS_ZEROGPU:
@@ -37,14 +45,13 @@ if IS_ZEROGPU:
37
  elif IS_T4_GPU:
38
  DEFAULT_BATCH_SIZE = 64
39
  DEFAULT_NUM_WORKERS = 4
40
- GPU_TYPE = "T4 Small"
41
  else:
42
  DEFAULT_BATCH_SIZE = 64
43
  DEFAULT_NUM_WORKERS = 8
44
- GPU_TYPE = "Standard GPU"
45
 
46
  import pickle
47
- import torch
48
  import pandas as pd
49
  import gradio as gr
50
  from pathlib import Path
@@ -53,10 +60,11 @@ from mussel.utils import get_features, segment_tissue, filter_features
53
  from mussel.utils.segment import draw_slide_mask
54
  from mussel.cli.tessellate import BiopsySegConfig, ResectionSegConfig, TcgaSegConfig
55
  from loguru import logger
56
-
57
-
58
  from mosaic.inference import run_aeon, run_paladin
59
 
 
 
 
60
 
61
  def _extract_ctranspath_features(coords, slide_path, attrs, num_workers):
62
  """Extract CTransPath features on GPU.
 
26
  return lambda f: f
27
  return fn
28
 
29
+ # Detect T4 hardware by checking actual GPU
30
+ import torch
31
+ IS_T4_GPU = False
32
+ GPU_NAME = "Unknown"
33
+ if not IS_ZEROGPU and torch.cuda.is_available():
34
+ try:
35
+ GPU_NAME = torch.cuda.get_device_name(0)
36
+ IS_T4_GPU = "T4" in GPU_NAME
37
+ except:
38
+ pass
39
 
40
  # Set optimal parameters based on hardware
41
  if IS_ZEROGPU:
 
45
  elif IS_T4_GPU:
46
  DEFAULT_BATCH_SIZE = 64
47
  DEFAULT_NUM_WORKERS = 4
48
+ GPU_TYPE = f"T4 ({GPU_NAME})"
49
  else:
50
  DEFAULT_BATCH_SIZE = 64
51
  DEFAULT_NUM_WORKERS = 8
52
+ GPU_TYPE = f"Standard GPU ({GPU_NAME})"
53
 
54
  import pickle
 
55
  import pandas as pd
56
  import gradio as gr
57
  from pathlib import Path
 
60
  from mussel.utils.segment import draw_slide_mask
61
  from mussel.cli.tessellate import BiopsySegConfig, ResectionSegConfig, TcgaSegConfig
62
  from loguru import logger
 
 
63
  from mosaic.inference import run_aeon, run_paladin
64
 
65
+ # Log hardware detection at module load
66
+ logger.info(f"Hardware: {GPU_TYPE} | batch_size={DEFAULT_BATCH_SIZE}, num_workers={DEFAULT_NUM_WORKERS}")
67
+
68
 
69
  def _extract_ctranspath_features(coords, slide_path, attrs, num_workers):
70
  """Extract CTransPath features on GPU.