ato 2.0.4__py3-none-any.whl

ato/db_routers/sql/manager.py

@@ -0,0 +1,188 @@
+import datetime
+from typing import Optional
+
+from sqlalchemy import create_engine, func
+from sqlalchemy.orm import sessionmaker, Session as SessionType
+
+from ato.db_routers import BaseLogger, BaseFinder
+from ato.db_routers.sql.schema import Base, Project, Experiment, Metric, Artifact, Fingerprint
+
+
+class SQLLogger(BaseLogger):
+    registry: set[str] = set()
+
+    def __init__(self, config):
+        super().__init__(config)
+        db_path = self.config.experiment.sql.db_path
+        self.engine = create_engine(db_path)
+        registry = self.__class__.registry
+        if db_path not in registry:
+            Base.metadata.create_all(self.engine)  # Base.metadata is SQLAlchemy's internal attribute
+            registry.add(db_path)
+        self.session = sessionmaker(bind=self.engine)()
+        self._current_run_id = None
+
+    def get_current_run(self):
+        return self.session.get(Experiment, self.get_current_run_id())
+
+    def get_current_run_id(self):
+        return self._current_run_id
+
+    def get_or_create_project(self):
+        project = self.session.query(Project).filter_by(name=self.config.experiment.project_name).first()
+        if not project:
+            project = Project(name=self.config.experiment.project_name)
+            self.session.add(project)
+            self.session.commit()
+        return project
+
+    def update_status(self, status):
+        run = self.get_current_run()
+        if run:
+            run.status = status
+            self.session.commit()
+
+    def run(self, tags=None):
+        project = self.get_or_create_project()
+        structural_hash = self.config.get_structural_hash()
+        run = Experiment(
+            project_id=project.id,
+            config=self.config.to_dict(),  # Recursively convert ADict to dict for JSON serialization
+            structural_hash=structural_hash,
+            tags=tags or []
+        )
+        self.session.add(run)
+        self.session.commit()
+        self._current_run_id = run.id
+        return run.id
+
+    def log_metric(self, key, value, step):
+        metric = Metric(
+            run_id=self.get_current_run_id(),
+            key=key,
+            value=value,
+            step=step
+        )
+        self.session.add(metric)
+        self.session.commit()
+
+    def log_artifact(self, run_id, file_path, data_type, metadata=None):
+        artifact = Artifact(
+            run_id=run_id,
+            path=file_path,
+            data_type=data_type,
+            data_info=metadata  # Column name is data_info in schema
+        )
+        self.session.add(artifact)
+        self.session.commit()
+
+    def finish(self, status='completed'):
+        run = self.get_current_run()
+        if run:
+            run.status = status
+            run.end_time = datetime.datetime.now(datetime.timezone.utc)
+            self.session.commit()
+
+
+class SQLFinder(BaseFinder):
+    def __init__(self, config):
+        super().__init__(config)
+        self.engine = create_engine(self.config.experiment.sql.db_path)
+        self.session_maker = sessionmaker(bind=self.engine)
+
+    def _get_session(self) -> SessionType:
+        return self.session_maker()
+
+    def find_project(self, project_name) -> Optional[Project]:
+        with self._get_session() as session:
+            return session.query(Project).filter_by(name=project_name).first()
+
+    def find_run(self, run_id: int) -> Optional[Experiment]:
+        with self._get_session() as session:
+            return session.get(Experiment, run_id)
+
+    def get_runs_in_project(self, project_name) -> list[Experiment]:
+        project = self.find_project(project_name)
+        if not project:
+            return []
+        else:
+            with self._get_session() as session:
+                return session.query(Experiment).filter_by(project_id=project.id).all()
+
+    def find_similar_runs(self, run_id: int) -> list[Experiment]:
+        with self._get_session() as session:
+            base_run = session.query(Experiment.structural_hash).filter_by(id=run_id).first()
+            if not base_run:
+                return []
+            else:
+                target_hash = base_run.structural_hash
+                return session.query(Experiment).filter(
+                    Experiment.structural_hash == target_hash,
+                    Experiment.id != run_id
+                ).all()
+
+    def find_similar_runs_by_trace(
+        self,
+        run_id: int,
+        trace_id: str,
+        trace_type: str = 'static'
+    ) -> list[Experiment]:
+        with self._get_session() as session:
+            target_fingerprint = session.query(Fingerprint.fingerprint).filter(
+                Fingerprint.run_id == run_id,
+                Fingerprint.trace_id == trace_id,
+                Fingerprint.trace_type == trace_type
+            ).scalar()
+            if not target_fingerprint:
+                return []
+            query = session.query(Fingerprint.run_id).filter(
+                Fingerprint.trace_id == trace_id,
+                Fingerprint.trace_type == trace_type,
+                Fingerprint.fingerprint == target_fingerprint,
+                Fingerprint.run_id != run_id
+            ).distinct()
+            return session.query(Experiment).filter(Experiment.id.in_(query)).all()
+
+    def find_best_run(self, project_name: str, metric_key: str, mode: str = 'max') -> Experiment | dict:
+        project = self.find_project(project_name)
+        if project:
+            with self._get_session() as session:
+                order_by_col = Metric.value.desc() if mode == 'max' else Metric.value.asc()
+                best_run = session.query(Experiment, Metric.value).join(
+                    Metric,
+                    Experiment.id == Metric.run_id
+                ).filter(
+                    Experiment.project_id == project.id,
+                    Metric.key == metric_key
+                ).order_by(
+                    order_by_col
+                ).first()
+                if best_run:
+                    return best_run[0]
+                else:
+                    return {'error': 'Run not found'}
+        return {'error': 'Project not found'}
+
+    def get_trace_statistics(self, project_name: str, trace_id: str) -> dict:
+        project = self.find_project(project_name)
+        if not project:
+            return {'error': 'Project not found'}
+        else:
+            with self._get_session() as session:
+                run_ids_in_project = session.query(Experiment.id).filter_by(project_id=project.id)
+                static_count = session.query(func.count(Fingerprint.fingerprint.distinct())).filter(
+                    Fingerprint.run_id.in_(run_ids_in_project),
+                    Fingerprint.trace_id == trace_id,
+                    Fingerprint.trace_type == 'static'
+                ).scalar()
+                runtime_count = session.query(func.count(Fingerprint.fingerprint.distinct())).filter(
+                    Fingerprint.run_id.in_(run_ids_in_project),
+                    Fingerprint.trace_id == trace_id,
+                    Fingerprint.trace_type == 'runtime'  # @runtime_trace의 타입
+                ).scalar()
+                return {
+                    'project_name': project_name,
+                    'trace_id': trace_id,
+                    'static_trace_versions': static_count,
+                    'runtime_trace_versions': runtime_count
+                }

ato/db_routers/sql/schema.py

@@ -0,0 +1,83 @@
+import datetime
+
+from sqlalchemy import Column, Integer, String, Float, DateTime, ForeignKey, JSON, Text
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import relationship
+
+Base = declarative_base()
+
+
+# single project
+class Project(Base):
+    __tablename__ = 'projects'
+
+    id = Column(Integer, primary_key=True)
+    name = Column(String, unique=True, nullable=False, index=True)
+    description = Column(Text)
+    created_at = Column(DateTime, default=datetime.datetime.utcnow)
+
+    runs = relationship('Experiment', back_populates='project')  # makes 1:N relationship
+
+
+# single experiment
+class Experiment(Base):
+    __tablename__ = 'experiments'
+
+    id = Column(Integer, primary_key=True)
+    project_id = Column(Integer, ForeignKey('projects.id'), nullable=False, index=True)
+    config = Column(JSON, nullable=False)
+
+    structural_hash = Column(String(64), index=True)
+
+    tags = Column(JSON, default=[])
+    status = Column(String, default='running', index=True)  # current status of experiment
+    start_time = Column(DateTime, default=datetime.datetime.utcnow)
+    end_time = Column(DateTime)
+
+    project = relationship('Project', back_populates='runs')
+    metrics = relationship('Metric', back_populates='run')
+    artifacts = relationship('Artifact', back_populates='run')
+    fingerprints = relationship('Fingerprint', back_populates='run')
+
+
+# single metric
+class Metric(Base):
+    __tablename__ = 'metrics'
+
+    id = Column(Integer, primary_key=True)
+    run_id = Column(Integer, ForeignKey('experiments.id'), nullable=False, index=True)
+
+    key = Column(String, nullable=False, index=True)  # name of metric
+    value = Column(Float, nullable=False)  # value of metric
+    step = Column(Integer, nullable=False)  # logged step
+    timestamp = Column(DateTime, default=datetime.datetime.now(datetime.timezone.utc))
+
+    run = relationship('Experiment', back_populates='metrics')
+
+
+# single artifact
+class Artifact(Base):
+    __tablename__ = 'artifacts'
+
+    id = Column(Integer, primary_key=True)
+    run_id = Column(Integer, ForeignKey('experiments.id'), nullable=False, index=True)
+
+    path = Column(String, nullable=False)
+    data_type = Column(String)
+    data_info = Column(JSON)
+
+    run = relationship('Experiment', back_populates='artifacts')
+
+
+# single fingerprint
+class Fingerprint(Base):
+    __tablename__ = 'fingerprints'
+
+    id = Column(Integer, primary_key=True)
+    run_id = Column(Integer, ForeignKey('experiments.id'), nullable=False, index=True)
+    trace_id = Column(String, nullable=False, index=True)
+
+    trace_type = Column(String(32), nullable=False, index=True)  # 'static' or 'runtime'
+    fingerprint = Column(String(128), nullable=False, index=True)  # hash value
+
+    run = relationship('Experiment', back_populates='fingerprints')

ato/hyperopt/base.py

@@ -0,0 +1,144 @@
+import math
+import uuid
+from copy import deepcopy as dcp
+from itertools import product
+
+import numpy as np
+import torch.distributed as dist
+
+from ato.adict import ADict
+
+
+class HyperOpt:
+    def __init__(self, scope, search_spaces, tracker=None, mode='max'):
+        if mode not in ('min', 'max'):
+            raise ValueError('mode must be either "min" or "max".')
+        self.scope = scope
+        self.search_spaces = search_spaces
+        self.config = scope.config.clone()
+        self.tracker = tracker
+        self.mode = mode
+        self.config.__hyperopt_id__ = self.get_hyperopt_id()
+
+    @classmethod
+    def get_hyperopt_id(cls):
+        return str(uuid.uuid4())
+
+    def main(self, func):
+        raise NotImplementedError()
+
+
+class DistributedMixIn:
+    def __init__(self, rank=0, world_size=1, backend='pytorch'):
+        self.rank = rank
+        self.world_size = world_size
+        self.backend = backend
+
+    @property
+    def is_root(self):
+        return self.rank == 0
+
+    def broadcast_object_from_root(self, obj):
+        if self.backend == 'pytorch':
+            obj = [obj]
+            dist.broadcast_object_list(obj)
+            obj = obj[0]
+        else:
+            raise ValueError(f'Unsupported backend: {self.backend}')
+        return obj
+
+    def all_gather_object(self, obj):
+        if self.backend == 'pytorch':
+            gathered_objects = [None for _ in range(self.world_size)]
+            dist.all_gather_object(gathered_objects, obj)
+        else:
+            raise ValueError(f'Unsupported backend: {self.backend}')
+        return gathered_objects
+
+    def destroy(self):
+        if self.backend == 'pytorch':
+            if dist.is_initialized():
+                dist.destroy_process_group()
+        else:
+            raise ValueError(f'Unsupported backend: {self.backend}')
+
+    def get_hyperopt_id(self):
+        return self.broadcast_object_from_root(str(uuid.uuid4()))
+
+
+class DistributedHyperOpt(DistributedMixIn, HyperOpt):
+    def __init__(
+        self,
+        scope,
+        search_spaces,
+        tracker=None,
+        mode='max',
+        rank=0,
+        world_size=1,
+        backend='pytorch'
+    ):
+        HyperOpt.__init__(self, scope, search_spaces, tracker, mode)
+        DistributedMixIn.__init__(self, rank, world_size, backend)
+
+
+class GridSpaceMixIn:
+    @classmethod
+    def prepare_distributions(cls, base_config, search_spaces):
+        sampling_spaces = ADict()
+        for param_name, search_space in search_spaces.items():
+            if 'param_type' not in search_space:
+                raise KeyError(f'param_type for parameter {param_name} is not defined at search_spaces.')
+            param_type = search_space['param_type'].upper()
+            if param_type == 'INTEGER':
+                start, stop = search_space.param_range
+                space_type = search_space.get('space_type', 'LINEAR')
+                if space_type == 'LINEAR':
+                    optim_space = np.linspace(
+                        start=start,
+                        stop=stop,
+                        num=search_space.num_samples,
+                        dtype=np.int64
+                    ).tolist()
+                elif space_type == 'LOG':
+                    base = search_space.get('base', 2)
+                    optim_space = np.logspace(
+                        start=math.log(start, base),
+                        stop=math.log(stop, base),
+                        num=search_space.num_samples,
+                        dtype=np.int64,
+                        base=base
+                    ).tolist()
+                else:
+                    raise ValueError(f'Invalid space_type: {space_type}')
+            elif param_type == 'FLOAT':
+                start, stop = search_space.param_range
+                space_type = search_space.get('space_type', 'LINEAR')
+                if space_type == 'LINEAR':
+                    optim_space = np.linspace(
+                        start=start,
+                        stop=stop,
+                        num=search_space.num_samples,
+                        dtype=np.float32
+                    ).tolist()
+                elif space_type == 'LOG':
+                    base = search_space.get('base', 10)
+                    optim_space = np.logspace(
+                        start=math.log(start, base),
+                        stop=math.log(stop, base),
+                        num=search_space.num_samples,
+                        dtype=np.float32,
+                        base=base
+                    ).tolist()
+                else:
+                    raise ValueError(f'Invalid space_type: {space_type}')
+            elif param_type == 'CATEGORY':
+                optim_space = search_space.categories
+            else:
+                raise ValueError(f'Unknown param_type for parameter {param_name}; {param_type}')
+            sampling_spaces[param_name] = optim_space
+        grid_space = [ADict(zip(sampling_spaces.keys(), values)) for values in product(*sampling_spaces.values())]
+        distributions = [
+            dcp(base_config).update(**partial_config)
+            for index, partial_config in enumerate(grid_space)
+        ]
+        return distributions

ato/hyperopt/hyperband.py

@@ -0,0 +1,103 @@
+import math
+from copy import deepcopy as dcp
+from itertools import chain
+
+from ato.adict import ADict
+from ato.hyperopt.base import HyperOpt, DistributedMixIn, GridSpaceMixIn
+
+
+class HyperBand(HyperOpt, GridSpaceMixIn):
+    def __init__(self, scope, search_spaces, halving_rate, num_min_samples, tracker=None, mode='max'):
+        if halving_rate <= 0 or halving_rate >= 1:
+            raise ValueError(f'halving_rate must be greater than 0.0 but less than 1.0, but got {halving_rate}.')
+        if num_min_samples < 1:
+            raise ValueError(f'num_min_samples must be greater than or equal to 1, but got {num_min_samples}.')
+        super().__init__(scope, search_spaces, tracker, mode)
+        self.halving_rate = halving_rate
+        self.num_min_samples = num_min_samples
+        self.distributions = self.prepare_distributions(self.config, self.search_spaces)
+
+    @classmethod
+    def prepare_distributions(cls, base_config, search_spaces):
+        distributions = super().prepare_distributions(base_config, search_spaces)
+        for distribution in distributions:
+            distribution.__num_halved__ = 0
+        return distributions
+
+    def main(self, func):
+        def launch(*args, **kwargs):
+            logs = []
+            distributions = self.distributions
+            while len(distributions) >= self.num_min_samples:
+                results = self.estimate(func, distributions, *args, **kwargs)
+                results.sort(key=lambda item: item.__metric__, reverse=self.mode == 'max')
+                logs.append(results)
+                distributions = []
+                for config in results[:int(len(results)*self.halving_rate)]:
+                    config.__num_halved__ += 1
+                    distributions.append(config)
+            last_config = logs[-1][0]
+            metric = self.estimate_single_run(func, last_config, *args, **kwargs)
+            best_config = dcp(last_config)
+            best_config.__metric__ = metric
+            logs.append([best_config])
+            return ADict(config=best_config, metric=metric, logs=logs)
+        return launch
+
+    def estimate(self, estimator, distributions, *args, **kwargs):
+        results = []
+        for config in distributions:
+            config = dcp(config)
+            self.scope.config = config
+            metric = self.estimate_single_run(estimator, config, *args, **kwargs)
+            config.__metric__ = metric
+            results.append(config)
+        return results
+
+    def estimate_single_run(self, estimator, config, *args, **kwargs):
+        self.scope.config = config
+        return self.scope(estimator)(*args, **kwargs)
+
+    def num_generations(self):
+        max_size = len(self.distributions)
+        min_size = self.num_min_samples
+        return math.ceil(math.log(min_size/max_size)/math.log(self.halving_rate))
+
+    def compute_optimized_initial_training_steps(self, max_steps):
+        num_generations = self.num_generations()
+        min_steps = max_steps*math.pow(self.halving_rate, num_generations)
+        return [
+            *(math.ceil(min_steps/math.pow(self.halving_rate, index)) for index in range(1, num_generations)),
+            max_steps
+        ]
+
+
+class DistributedHyperBand(DistributedMixIn, HyperBand, GridSpaceMixIn):
+    def __init__(
+        self,
+        scope,
+        search_spaces,
+        halving_rate,
+        num_min_samples,
+        tracker=None,
+        mode='max',
+        rank=0,
+        world_size=1,
+        backend='pytorch'
+    ):
+        DistributedMixIn.__init__(self, rank, world_size, backend)
+        HyperBand.__init__(self, scope, search_spaces, halving_rate, num_min_samples, tracker, mode)
+
+    def estimate(self, estimator, distributions, *args, **kwargs):
+        batch_size = math.ceil(len(distributions)/self.world_size)
+        distributions = distributions[self.rank*batch_size:(self.rank+1)*batch_size]
+        results = super().estimate(estimator, distributions, *args, **kwargs)
+        results = list(chain(*self.all_gather_object(results)))
+        return results
+
+    def estimate_single_run(self, estimator, config, *args, **kwargs):
+        result = super().estimate(estimator, config, *args, **kwargs)
+        self.destroy()
+        return result
+
+

ato/parser.py

@@ -0,0 +1,103 @@
+def parse_command(command):
+    tokens = []
+    i = 0
+    length = len(command)
+    while i < length:
+        while i < length and command[i].isspace():
+            i += 1
+        if i >= length:
+            break
+        start = i
+        while i < length and not command[i].isspace() and command[i] != '=':
+            i += 1
+        if i < length and command[i] == '=':
+            key = command[start:i]
+            i += 1
+            if i < length:
+                value, i = parse_value(command, i)
+            else:
+                value = ''
+            tokens.append(f'{key}={value}')
+        else:
+            token_start = start
+            while i < length and not command[i].isspace():
+                i += 1
+            tokens.append(command[token_start:i])
+    return tokens
+
+
+def parse_value(command, i):
+    if i < len(command):
+        if command[i] == '%':
+            return parse_backtick_string(command, i)
+        elif command[i] in ['[', '(', '{']:
+            return parse_bracketed_value(command, i)
+        else:
+            start = i
+            while i < len(command) and not command[i].isspace():
+                i += 1
+            return command[start:i], i
+    return '', i
+
+
+def parse_backtick_string(command, i):
+    assert command[i] == '%'
+    i += 1
+    value = ['%']
+    length = len(command)
+    nesting_level = 1
+    while i < length:
+        c = command[i]
+        if c == '\\' and i+1 < length:
+            value.append(c)
+            value.append(command[i+1])
+            i += 2
+        elif c == '%':
+            value.append(c)
+            i += 1
+            if i < length and command[i] == '%':
+                value.append(command[i])
+                i += 1
+            else:
+                nesting_level -= 1
+                if nesting_level == 0:
+                    break
+        elif c == '%':
+            value.append(c)
+            nesting_level += 1
+            i += 1
+        else:
+            value.append(c)
+            i += 1
+    return ''.join(value), i
+
+
+def parse_bracketed_value(command, i):
+    brackets = {'[': ']', '{': '}', '(': ')'}
+    opening_bracket = command[i]
+    closing_bracket = brackets[opening_bracket]
+    value = [opening_bracket]
+    i += 1
+    length = len(command)
+    stack = [closing_bracket]
+    while i < length and stack:
+        c = command[i]
+        if c == '%':
+            backtick_value, i = parse_backtick_string(command, i)
+            value.append(backtick_value)
+        elif c == '\\' and i+1 < length:
+            value.append(c)
+            value.append(command[i+1])
+            i += 2
+        elif c in brackets:
+            stack.append(brackets[c])
+            value.append(c)
+            i += 1
+        elif c == stack[-1]:
+            stack.pop()
+            value.append(c)
+            i += 1
+        else:
+            value.append(c)
+            i += 1
+    return ''.join(value), i