gary-boon Claude Opus 4.5 commited on
Commit
6435a75
Β·
1 Parent(s): d129e37

Make zarr/numcodecs imports optional for ARM64 compatibility

Browse files

- Use try/except for zarr imports
- Add ZARR_AVAILABLE flag
- Skip storage operations gracefully when zarr not available
- Core model inference still works without tensor caching

πŸ€– Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Files changed (1) hide show
  1. backend/storage.py +27 -8
backend/storage.py CHANGED
@@ -16,8 +16,16 @@ Storage structure:
16
  └── telemetry.jsonl
17
  """
18
 
19
- import zarr
20
- import numcodecs
 
 
 
 
 
 
 
 
21
  import numpy as np
22
  import torch
23
  import json
@@ -53,12 +61,16 @@ class ZarrStorage:
53
  # Create directories
54
  self.tensor_dir.mkdir(parents=True, exist_ok=True)
55
 
56
- # Blosc compressor for efficient compression
57
- self.compressor = numcodecs.Blosc(
58
- cname='zstd', # zstd algorithm (good compression + speed)
59
- clevel=5, # Compression level (1-9, 5 is balanced)
60
- shuffle=numcodecs.Blosc.SHUFFLE # Byte shuffle for better compression
61
- )
 
 
 
 
62
 
63
  def save_instrumentation_data(self, data: InstrumentationData) -> Dict[str, Any]:
64
  """
@@ -70,6 +82,10 @@ class ZarrStorage:
70
  Returns:
71
  Dictionary with file paths and sizes
72
  """
 
 
 
 
73
  logger.info(f"Saving instrumentation data for run {self.run_id}...")
74
 
75
  result = {
@@ -227,6 +243,9 @@ class ZarrStorage:
227
  Returns:
228
  Attention matrix [seq_len, seq_len]
229
  """
 
 
 
230
  path = self.tensor_dir / "attention.zarr"
231
 
232
  if not path.exists():
 
16
  └── telemetry.jsonl
17
  """
18
 
19
+ # Optional imports for zarr storage (not available on ARM64)
20
+ try:
21
+ import zarr
22
+ import numcodecs
23
+ ZARR_AVAILABLE = True
24
+ except ImportError:
25
+ zarr = None
26
+ numcodecs = None
27
+ ZARR_AVAILABLE = False
28
+
29
  import numpy as np
30
  import torch
31
  import json
 
61
  # Create directories
62
  self.tensor_dir.mkdir(parents=True, exist_ok=True)
63
 
64
+ # Blosc compressor for efficient compression (if zarr available)
65
+ if ZARR_AVAILABLE:
66
+ self.compressor = numcodecs.Blosc(
67
+ cname='zstd', # zstd algorithm (good compression + speed)
68
+ clevel=5, # Compression level (1-9, 5 is balanced)
69
+ shuffle=numcodecs.Blosc.SHUFFLE # Byte shuffle for better compression
70
+ )
71
+ else:
72
+ self.compressor = None
73
+ logger.warning("Zarr not available - tensor storage will be disabled")
74
 
75
  def save_instrumentation_data(self, data: InstrumentationData) -> Dict[str, Any]:
76
  """
 
82
  Returns:
83
  Dictionary with file paths and sizes
84
  """
85
+ if not ZARR_AVAILABLE:
86
+ logger.warning("Zarr not available - skipping instrumentation data save")
87
+ return {'run_id': self.run_id, 'error': 'zarr_not_available'}
88
+
89
  logger.info(f"Saving instrumentation data for run {self.run_id}...")
90
 
91
  result = {
 
243
  Returns:
244
  Attention matrix [seq_len, seq_len]
245
  """
246
+ if not ZARR_AVAILABLE:
247
+ raise RuntimeError("Zarr not available - cannot load attention data")
248
+
249
  path = self.tensor_dir / "attention.zarr"
250
 
251
  if not path.exists():