aiauto-client 0.1.0__py3-none-any.whl

aiauto/__init__.py

@@ -0,0 +1,50 @@
+from .core import AIAutoController, TrialController, CallbackTopNArtifact, StudyWrapper
+from .serialization import SourceCodeSerializer, create_study_with_source_serialization
+
+__version__ = "0.1.0"
+
+__all__ = [
+    'AIAutoController',
+    'TrialController', 
+    'CallbackTopNArtifact',
+    'StudyWrapper',
+    'SourceCodeSerializer',
+    'create_study_with_source_serialization',
+]
+
+# Optuna 호환성을 위한 간편 함수
+def create_study(
+    objective=None,
+    study_name='aiauto_study', 
+    direction='minimize',
+    **kwargs
+):
+    """
+    Optuna 호환 create_study 함수
+    
+    사용법:
+        study = aiauto.create_study(
+            objective=my_objective,
+            study_name='my_study',
+            direction='maximize'
+        )
+        study.optimize(n_trials=100)
+    """
+    controller = AIAutoController()
+    
+    if objective is not None:
+        return controller.create_study_with_serialization(
+            objective=objective,
+            study_name=study_name,
+            direction=direction,
+            **kwargs
+        )
+    else:
+        # objective가 없으면 일반 optuna study 반환 (기존 방식)
+        import optuna
+        return optuna.create_study(
+            study_name=study_name,
+            direction=direction,
+            storage=controller.get_storage(),
+            **kwargs
+        )

aiauto/core.py

@@ -0,0 +1,257 @@
+from os import makedirs, environ
+import tempfile
+from typing import Union, Callable, Dict, Any, Optional
+import optuna
+
+from .serialization import create_study_with_source_serialization, SourceCodeSerializer
+
+
+class AIAutoController:
+    # singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, "_instance"):
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    # singleton pattern
+    def __init__(self):
+        cls = type(self)
+        if not hasattr(cls, "_init"):
+            # singleton pattern
+            # ---------------------
+
+            # TODO token 인증
+            token = environ.get('AIAUTO_TOKEN')
+
+            # mode별 storage 설정
+            mode = environ.get('AIAUTO_MODE', 'single_gpu')
+            if mode == "distributed":
+                # DDP/FSDP pruning callback 지원을 위해 RDBStorage 사용
+                self.storage = optuna.storages.RDBStorage(
+                    url="sqlite:///optuna.db",
+                    engine_kwargs={"connect_args": {"timeout": 10}}
+                )
+            else:
+                # 기본 GrpcStorageProxy (single GPU 등)
+                self.storage = optuna.storages.GrpcStorageProxy(host="localhost", port=13000)
+
+            # artifact storage
+            # TODO 나중에 s3 던 다른 mount 된 경로 건 바꿔야 함
+            makedirs('./artifacts', exist_ok=True)
+            self.artifact_store = optuna.artifacts.FileSystemArtifactStore('./artifacts')
+            # model 저장을 위한 임시 디렉토리
+            self.tmp_dir = tempfile.mkdtemp(prefix=f'ai_auto_tmp_')
+
+            # ---------------------
+            # singleton pattern end
+            cls._init = True
+
+    def get_storage(self):
+        return self.storage
+
+    def get_artifact_store(self) -> Union[
+        optuna.artifacts.FileSystemArtifactStore,
+        optuna.artifacts.Boto3ArtifactStore,
+        optuna.artifacts.GCSArtifactStore,
+    ]:
+        return self.artifact_store
+
+    def get_artifact_tmp_dir(self):
+        return self.tmp_dir
+
+    def create_study_with_serialization(
+        self,
+        objective: Callable,
+        study_name: str,
+        direction: str = 'minimize',
+        sampler: Optional[optuna.samplers.BaseSampler] = None,
+        pruner: Optional[optuna.pruners.BasePruner] = None,
+        **optuna_kwargs
+    ) -> 'StudyWrapper':
+        """
+        소스코드 직렬화를 사용하여 Study 생성
+        
+        Args:
+            objective: HPO에 사용할 objective 함수
+            study_name: Study 이름
+            direction: 최적화 방향 ('minimize' 또는 'maximize')
+            sampler: Optuna sampler (기본값: TPESampler)
+            pruner: Optuna pruner
+            **optuna_kwargs: optuna.create_study에 전달할 추가 인자
+            
+        Returns:
+            StudyWrapper 객체 (Optuna Study 호환)
+        """
+        study_config = {
+            'study_name': study_name,
+            'direction': direction,
+            'sampler': sampler.__class__.__name__ if sampler else 'TPESampler',
+            'pruner': pruner.__class__.__name__ if pruner else None,
+        }
+        
+        # 소스코드 직렬화
+        serialized_objective, processed_config = create_study_with_source_serialization(
+            objective, study_config, **optuna_kwargs
+        )
+        
+        # StudyWrapper 생성 (실제 gRPC 전송은 optimize 시점에)
+        return StudyWrapper(
+            serialized_objective=serialized_objective,
+            study_config=processed_config,
+            storage=self.storage,
+            artifact_store=self.artifact_store
+        )
+
+
+class TrialController:
+    def __init__(self, trial: optuna.trial.Trial):
+        self.trial = trial
+        self.logger = optuna.logging.get_logger("optuna")
+        self.logs = []
+
+    def get_trial(self) -> optuna.trial.Trial:
+        return self.trial
+
+    def log(self, value: str):
+        # optuna dashboard 에 log 를 확인하는 기능이 없어서 user_attribute 에 log를 확인할 수 있게 추가
+        self.logs.append(value)
+        self.trial.set_user_attr('logs', ' '.join([f"[{i+1:05d}] {log}" for i, log in enumerate(self.logs)]))
+        # 실제 log 를 trial_number 랑 같이 확인할 수 있게
+        self.logger.info(f'\ntrial_number: {self.trial.number}, {value}')
+
+
+# 용량 제한으로 상위 N개의 trial artifact 만 유지
+class CallbackTopNArtifact:
+    def __init__(
+        self,
+        artifact_store: Union[
+            optuna.artifacts.FileSystemArtifactStore,
+            optuna.artifacts.Boto3ArtifactStore,
+            optuna.artifacts.GCSArtifactStore,
+        ],
+        artifact_attr_name: str = 'artifact_id',
+        n_keep: int = 5,
+    ):
+        self.artifact_store = artifact_store
+        self.check_attr_name = artifact_attr_name
+        self.n_keep = n_keep
+
+    def __call__(self, study: optuna.study.Study, trial: optuna.trial.FrozenTrial):
+        # COMPLETE 상태이고 artifact를 가진 trial들만 정렬
+        finished_with_artifacts = [
+            t for t in study.trials
+            if t.state == optuna.trial.TrialState.COMPLETE and self.check_attr_name in t.user_attrs
+        ]
+
+        # 방향에 따라 정렬 (maximize면 내림차순, minimize면 오름차순)
+        reverse_sort = study.direction == optuna.study.StudyDirection.MAXIMIZE
+        finished_with_artifacts.sort(key=lambda t: t.value, reverse=reverse_sort)
+
+        # 상위 n_keep개 초과하는 trial들의 artifact 삭제
+        for old_trial in finished_with_artifacts[self.n_keep:]:
+            artifact_id = old_trial.user_attrs.get(self.check_attr_name)
+            if artifact_id:
+                try:
+                    self.artifact_store.remove(artifact_id)
+                    # user_attr에서도 제거
+                    study._storage.set_trial_user_attr(old_trial._trial_id, self.check_attr_name, None)
+                except Exception as e:
+                    print(f"Warning: Failed to remove artifact {artifact_id}: {e}")
+
+
+class StudyWrapper:
+    """
+    Optuna Study 호환성을 제공하는 래퍼 클래스
+    
+    이 클래스는 소스코드 직렬화된 objective 함수를 관리하고
+    실제 HPO 실행을 위해 gRPC 백엔드와 통신합니다.
+    """
+    
+    def __init__(
+        self,
+        serialized_objective: Dict[str, Any],
+        study_config: Dict[str, Any],
+        storage,
+        artifact_store
+    ):
+        self.serialized_objective = serialized_objective
+        self.study_config = study_config
+        self.storage = storage
+        self.artifact_store = artifact_store
+        self._local_study = None  # 로컬 테스트용
+        
+    def optimize(
+        self,
+        n_trials: int = 100,
+        n_jobs: int = 1,
+        callbacks: Optional[list] = None,
+        **kwargs
+    ):
+        """
+        HPO 최적화 실행
+        
+        실제 구현에서는 gRPC를 통해 백엔드로 전송하지만,
+        현재는 로컬에서 역직렬화하여 테스트합니다.
+        """
+        print("🚀 Starting HPO optimization with source code serialization...")
+        print(f"📊 Study: {self.study_config['study_name']}")
+        print(f"🎯 Direction: {self.study_config['direction']}")
+        print(f"🔢 Trials: {n_trials}")
+        
+        try:
+            # 소스코드 역직렬화로 objective 함수 복원
+            objective_func = SourceCodeSerializer.deserialize_objective(
+                self.serialized_objective
+            )
+            print("✅ Objective function deserialized successfully")
+            
+            # 로컬 Study 생성 (실제로는 gRPC 통신)
+            self._local_study = optuna.create_study(
+                study_name=self.study_config['study_name'],
+                direction=self.study_config['direction'],
+                storage=self.storage,
+                load_if_exists=True
+            )
+            
+            # 최적화 실행
+            self._local_study.optimize(
+                objective_func,
+                n_trials=n_trials,
+                n_jobs=n_jobs,
+                callbacks=callbacks or [],
+                **kwargs
+            )
+            
+            print(f"🎉 Optimization completed! Best value: {self.best_value}")
+            
+        except Exception as e:
+            print(f"❌ Optimization failed: {e}")
+            raise
+    
+    @property
+    def best_trial(self):
+        """최고 성능 Trial 반환"""
+        if self._local_study:
+            return self._local_study.best_trial
+        return None
+    
+    @property
+    def best_value(self):
+        """최고 성능 값 반환"""
+        if self._local_study:
+            return self._local_study.best_value
+        return None
+    
+    @property
+    def best_params(self):
+        """최고 성능 하이퍼파라미터 반환"""
+        if self._local_study:
+            return self._local_study.best_params
+        return None
+    
+    @property
+    def trials(self):
+        """모든 Trial 목록 반환"""
+        if self._local_study:
+            return self._local_study.trials
+        return []

aiauto/serialization.py

@@ -0,0 +1,138 @@
+"""
+Source Code Serialization Module
+
+이 모듈은 Python 버전 간 호환성을 위해 CloudPickle 대신 
+inspect.getsource를 사용한 소스코드 직렬화 방식을 제공합니다.
+"""
+
+import inspect
+import types
+from typing import Callable, Dict, Any, Tuple
+
+
+class SourceCodeSerializer:
+    """Objective 함수를 소스코드로 직렬화하는 클래스"""
+    
+    @staticmethod
+    def serialize_objective(objective_func: Callable) -> Dict[str, Any]:
+        """
+        Objective 함수를 소스코드로 직렬화
+        
+        Args:
+            objective_func: 직렬화할 objective 함수
+            
+        Returns:
+            직렬화된 데이터 딕셔너리
+            - source_code: 함수의 소스코드 문자열
+            - func_name: 함수 이름
+            - dependencies: 필요한 import 구문들
+        """
+        try:
+            # 함수 소스코드 추출
+            source_code = inspect.getsource(objective_func)
+            func_name = objective_func.__name__
+            
+            # 함수가 정의된 모듈의 정보 추출
+            module = inspect.getmodule(objective_func)
+            dependencies = []
+            
+            if module and hasattr(module, '__file__'):
+                # 모듈에서 import 구문들 추출 (간단한 방식)
+                with open(module.__file__, 'r') as f:
+                    module_source = f.read()
+                
+                # import 구문 추출 (개선된 파싱 필요시 ast 모듈 사용)
+                lines = module_source.split('\n')
+                for line in lines:
+                    line = line.strip()
+                    if line.startswith('import ') or line.startswith('from '):
+                        # 기본적인 import 구문만 추출
+                        if not any(skip in line for skip in ['client', '__', 'relative']):
+                            dependencies.append(line)
+            
+            return {
+                'source_code': source_code,
+                'func_name': func_name,
+                'dependencies': dependencies,
+                'serialization_method': 'source_code'
+            }
+            
+        except Exception as e:
+            raise RuntimeError(f"Failed to serialize objective function: {e}")
+    
+    @staticmethod
+    def deserialize_objective(serialized_data: Dict[str, Any]) -> Callable:
+        """
+        직렬화된 데이터로부터 objective 함수를 복원
+        
+        Args:
+            serialized_data: serialize_objective에서 생성된 데이터
+            
+        Returns:
+            복원된 objective 함수
+        """
+        try:
+            source_code = serialized_data['source_code']
+            func_name = serialized_data['func_name']
+            dependencies = serialized_data.get('dependencies', [])
+            
+            # 실행 네임스페이스 생성
+            exec_namespace = {'__builtins__': __builtins__}
+            
+            # 의존성 import 실행
+            for dep in dependencies:
+                try:
+                    exec(dep, exec_namespace)
+                except Exception as import_error:
+                    # import 실패는 경고만 하고 계속 진행
+                    print(f"Warning: Failed to import dependency '{dep}': {import_error}")
+            
+            # 소스코드 실행
+            exec(source_code, exec_namespace)
+            
+            # 함수 객체 추출
+            if func_name not in exec_namespace:
+                raise NameError(f"Function '{func_name}' not found in executed namespace")
+            
+            objective_func = exec_namespace[func_name]
+            
+            if not callable(objective_func):
+                raise TypeError(f"'{func_name}' is not callable")
+            
+            return objective_func
+            
+        except Exception as e:
+            raise RuntimeError(f"Failed to deserialize objective function: {e}")
+
+
+def create_study_with_source_serialization(
+    objective: Callable,
+    study_config: Dict[str, Any],
+    **optuna_kwargs
+) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+    """
+    소스코드 직렬화를 사용하여 study 생성 준비
+    
+    Args:
+        objective: HPO에 사용할 objective 함수
+        study_config: study 설정 (name, direction, sampler, pruner 등)
+        **optuna_kwargs: optuna.create_study에 전달할 추가 인자들
+        
+    Returns:
+        Tuple[serialized_objective, study_config]
+        - serialized_objective: 직렬화된 objective 함수 데이터
+        - study_config: study 설정 데이터
+    """
+    # Objective 함수 직렬화
+    serialized_objective = SourceCodeSerializer.serialize_objective(objective)
+    
+    # Study 설정 정리
+    processed_config = {
+        'study_name': study_config.get('study_name', 'unnamed_study'),
+        'direction': study_config.get('direction', 'minimize'),
+        'sampler': study_config.get('sampler', 'TPESampler'),
+        'pruner': study_config.get('pruner', None),
+        'optuna_kwargs': optuna_kwargs
+    }
+    
+    return serialized_objective, processed_config

aiauto_client-0.1.0.dist-info/METADATA

@@ -0,0 +1,74 @@
+Metadata-Version: 2.1
+Name: aiauto-client
+Version: 0.1.0
+Summary: AI Auto HPO (Hyperparameter Optimization) Client Library
+Author-email: AIAuto Team <ainode@zeroone.ai>
+Project-URL: Homepage, https://aiauto.cloude.ainode.ai
+Project-URL: Repository, https://aiauto.cloude.ainode.ai
+Project-URL: Documentation, https://aiauto.cloude.ainode.ai
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+Requires-Dist: optuna>=3.0.0
+Requires-Dist: grpcio>=1.50.0
+Requires-Dist: grpcio-tools>=1.50.0
+Requires-Dist: protobuf>=4.0.0
+
+# AIAuto - Hyperparameter Optimization Client Library
+
+AIAuto는 Kubernetes 기반의 분산 HPO(Hyperparameter Optimization) 시스템을 위한 클라이언트 라이브러리입니다.
+사용자 python lib <-> Next.js 서버 사이 gRPC 통신 담당
+
+## lib build
+- pypi build, upload 종속성 다운로드 `pip install build twine`
+- build lib `python -m build --wheel --sdist`
+  - `aiauto-0.1.0.whl` 생성
+  - `aiauto-0.1.0.tar.gz` 생성
+  - `aiauto.egg-info` 생성
+- `twine upload --repository testpypi dist/*`
+- `twine upload dist/*`
+  - upload 시 pypi token 을 입력하라고 나옴, pypi 로그인 계정 설정가면 있다 
+
+## 설치
+- `pip install aiauto`
+
+## 빠른 시작
+```python
+import aiauto
+
+# 컨트롤러 초기화
+ac = aiauto.AIAutoController()
+
+# Objective 함수 정의
+def objective(trial):
+    tc = aiauto.TrialController(trial)
+    
+    # 하이퍼파라미터 샘플링
+    lr = trial.suggest_float('lr', 1e-5, 1e-1, log=True)
+    
+    # 모델 학습 및 평가 로직
+    # ...
+    tc.log(f'full dataset: train {len(dataset)}, test {len(dataset_test)}, batch_size {batch_size}')
+    
+    return accuracy
+
+# Study 생성 및 최적화 실행
+study = optuna.create_study(
+    study_name='my_optimization',
+    storage=ac.get_storage(),
+    direction='maximize'
+)
+
+study.optimize(objective, n_trials=100)
+```

aiauto_client-0.1.0.dist-info/RECORD

@@ -0,0 +1,7 @@
+aiauto/__init__.py,sha256=VvEM3L0NZGrHi3kHV_gSRf8X2baqLDPOSbArgd6LpaI,1353
+aiauto/core.py,sha256=GKCF24GA25QCu8n2q3YXnff4Sb3Dfx1yKFvE7QZ8108,9182
+aiauto/serialization.py,sha256=6Rb5k01hx7uXaLt1XmUrmn1KzMjxsYinzi4fjglc3jw,5137
+aiauto_client-0.1.0.dist-info/METADATA,sha256=inimyQ0HuHH8mWkYUOP1FpXez506z1i8_oi0UVBSwsE,2510
+aiauto_client-0.1.0.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
+aiauto_client-0.1.0.dist-info/top_level.txt,sha256=Sk2ctO9_Bf_tAPwq1x6Vfl6OuL29XzwMTO4F_KG6oJE,7
+aiauto_client-0.1.0.dist-info/RECORD,,

aiauto_client-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (75.3.2)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

aiauto_client-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+aiauto