Coverage for streamwise/container_config.py: 95%

1""" 

2Shared container deployment configuration for StreamWise. 

3 

4Central source of truth for container resource defaults, MIG profiles, 

5GPU type mappings, and related deployment constants. Both allocator_bridge 

6and streamwise.py import from here to avoid duplication. 

7""" 

8 

9from __future__ import annotations 

10 

11import sys 

12import os 

13from dataclasses import dataclass 

14from typing import Union 

15 

16_HERE = os.path.dirname(os.path.abspath(__file__)) 

17if _HERE not in sys.path: 

18 sys.path.insert(0, _HERE) 

19 

20_REPO_ROOT = os.path.dirname(_HERE) 

21if _REPO_ROOT not in sys.path: 

22 sys.path.insert(0, _REPO_ROOT) 

23 

24# model_provisioner import adds simulator/ to sys.path 

25import model_provisioner # noqa: E402, F401 

26 

27from sim_types import GPUType # noqa: E402 

28from sim_types import Model # noqa: E402 

29 

30 

31@dataclass(frozen=True) 

32class ContainerResourceSpec: 

33 """Default resource settings for a container deployment.""" 

34 cpu: int 

35 memory_gib: int 

36 ephemeral_storage_gib: int 

37 gpu: Union[int, str] = 0 

38 

39 

40# Default CPU/memory/storage and baseline GPU settings for key services. 

41# Keep in sync with deployment defaults in streamwise/templates/add_pod.html. 

42MODEL_TO_CONTAINER_NAME: dict[Model, str] = { 

43 Model.GEMMA: "gemma", 

44 Model.FLUX: "flux", 

45 Model.HF: "hunyuanframepackf1", 

46 Model.HF_VAE: "hunyuanframepackvae", 

47 Model.FT: "fantasytalking", 

48} 

49 

50MODEL_CONTAINER_RESOURCES: dict[Model, ContainerResourceSpec] = { 

51 Model.GEMMA: ContainerResourceSpec(cpu=16, memory_gib=192, ephemeral_storage_gib=64, gpu=2), 

52 Model.FLUX: ContainerResourceSpec(cpu=12, memory_gib=128, ephemeral_storage_gib=64, gpu=2), 

53 Model.HF: ContainerResourceSpec(cpu=24, memory_gib=128, ephemeral_storage_gib=64, gpu=2), 

54 Model.HF_VAE: ContainerResourceSpec(cpu=4, memory_gib=32, ephemeral_storage_gib=16, gpu=1), 

55 Model.FT: ContainerResourceSpec(cpu=12, memory_gib=192, ephemeral_storage_gib=64, gpu=2), 

56} 

57 

58CONTAINER_RESOURCES: dict[str, ContainerResourceSpec] = { 

59 MODEL_TO_CONTAINER_NAME[model]: spec 

60 for model, spec in MODEL_CONTAINER_RESOURCES.items() 

61} 

62CONTAINER_RESOURCES.update({ 

63 "realesrgan": ContainerResourceSpec(cpu=4, memory_gib=32, ephemeral_storage_gib=16, gpu="1g.10gb"), 

64 "kokoro": ContainerResourceSpec(cpu=2, memory_gib=8, ephemeral_storage_gib=16, gpu="1g.10gb"), 

65 "yolo": ContainerResourceSpec(cpu=4, memory_gib=8, ephemeral_storage_gib=16, gpu="1g.10gb"), 

66}) 

67 

68# GPU type string used by pod_manager (lowercase). 

69GPU_TYPE_TO_POD_STR: dict[GPUType, str] = { 

70 GPUType.A100: "a100", 

71 GPUType.H100: "h100", 

72 GPUType.H200: "h200", 

73 GPUType.GB200: "gb200", 

74} 

75 

76# MIG is only supported by pod_manager on these GPU types. 

77MIG_CAPABLE_GPU_TYPES: frozenset[GPUType] = frozenset({GPUType.A100, GPUType.H100}) 

78 

79# Containers that prefer a MIG slice when the selected GPU type supports MIG. 

80# When MIG is available on the cluster, these services use a MIG slice (shared GPU). 

81# When MIG is NOT available, they fall back to 1 full GPU each and the extra GPUs 

82# are counted against the budget (with a warning if exceeded). 

83MIG_CONTAINERS: dict[str, str] = { 

84 "kokoro": "1g.10gb", 

85 "yolo": "1g.10gb", 

86 "realesrgan": "1g.10gb", 

87} 

88 

89# Containers that are co-located with their parent model (sharing GPUs on the same server). 

90# The allocator counts their GPUs as part of the parent model's allocation, so they should 

91# deploy with gpu=0 to avoid double-counting. 

92COLOCATED_CONTAINERS: frozenset[str] = frozenset({"hunyuanframepackvae"}) 

93 

94# Whether MIG is actually configured on the cluster. 

95# When False, MIG_CONTAINERS entries fall back to full GPUs. 

96MIG_AVAILABLE: bool = False 

97 

98 

99def get_minimum_service_container_specs(max_gpus: int) -> dict[str, ContainerResourceSpec]: 

100 """Build container defaults for /api/service minimal deployment. 

101 

102 Negative `max_gpus` values are clamped to 0 to avoid invalid GPU assignments. 

103 """ 

104 validated_max_gpus = max(0, max_gpus) 

105 container_specs: dict[str, ContainerResourceSpec] = { 

106 "podcasttranscript": ContainerResourceSpec(cpu=1, memory_gib=4, ephemeral_storage_gib=16, gpu=0), 

107 "slidetranscript": ContainerResourceSpec(cpu=1, memory_gib=4, ephemeral_storage_gib=16, gpu=0), 

108 "fluxkontext": ContainerResourceSpec(cpu=12, memory_gib=128, ephemeral_storage_gib=64, gpu=1), 

109 "whisper": ContainerResourceSpec(cpu=2, memory_gib=8, ephemeral_storage_gib=16, gpu=1), 

110 } 

111 

112 for name, spec in CONTAINER_RESOURCES.items(): 

113 if name in MIG_CONTAINERS: 

114 assigned_gpu: Union[int, str] = MIG_CONTAINERS[name] 

115 elif name == "hunyuanframepackvae": 

116 assigned_gpu = 1 

117 else: 

118 assigned_gpu = min(2, validated_max_gpus) 

119 

120 container_specs[name] = ContainerResourceSpec( 

121 cpu=spec.cpu, 

122 memory_gib=spec.memory_gib, 

123 ephemeral_storage_gib=spec.ephemeral_storage_gib, 

124 gpu=assigned_gpu, 

125 ) 

126 

127 return container_specs