arbor-ai 0.1.4__py3-none-any.whl → 0.1.6__py3-none-any.whl

arbor/server/services/job_manager.py

@@ -1,76 +1,81 @@
 import uuid
-from enum import Enum
-import logging
 from datetime import datetime
+from typing import Literal
+
+from arbor.server.api.models.schemas import JobStatus
+from arbor.server.core.config import Settings
+
+
+class JobEvent:
+    def __init__(
+        self, level: Literal["info", "warning", "error"], message: str, data: dict = {}
+    ):
+        self.level = level
+        self.message = message
+        self.data = data
+
+        self.id = str(f"ftevent-{uuid.uuid4()}")
+        self.created_at = datetime.now()
+
+
+class JobCheckpoint:
+    def __init__(
+        self,
+        fine_tuned_model_checkpoint: str,
+        fine_tuning_job_id: str,
+        metrics: dict,
+        step_number: int,
+    ):
+        self.id = str(f"ftckpt-{uuid.uuid4()}")
+        self.fine_tuned_model_checkpoint = fine_tuned_model_checkpoint
+        self.fine_tuning_job_id = fine_tuning_job_id
+        self.metrics = metrics
+        self.step_number = step_number
+        self.created_at = datetime.now()
 
-# https://platform.openai.com/docs/api-reference/fine-tuning/object
-class JobStatus(Enum):
-    PENDING = "pending" # Not in OAI
-    VALIDATING_FILES = "validating_files"
-    QUEUED = "queued"
-    RUNNING = "running"
-    SUCCEEDED = "succeeded"
-    FAILED = "failed"
-    CANCELLED = "cancelled"
-
-class JobLogHandler(logging.Handler):
-  def __init__(self, job):
-    super().__init__()
-    self.job = job
-
-  def emit(self, record):
-    log_entry = {
-      'timestamp': datetime.fromtimestamp(record.created).isoformat(),
-      'level': record.levelname,
-      'message': record.getMessage()
-    }
-    self.job.logs.append(log_entry)
 
 class Job:
-  def __init__(self, id: str, status: JobStatus):
-    self.id = id
-    self.status = status
-    self.fine_tuned_model = None
-    self.logs = []
-    self.logger = None
-    self.log_handler = None
-
-  def setup_logger(self, logger_name: str = None) -> logging.Logger:
-    """Sets up logging for the job with a dedicated handler."""
-    if logger_name is None:
-      logger_name = f"job_{self.id}"
-
-    logger = logging.getLogger(logger_name)
-    logger.setLevel(logging.INFO)
-
-    # Create and setup handler if not already exists
-    if self.log_handler is None:
-      handler = JobLogHandler(self)
-      formatter = logging.Formatter('%(message)s')
-      handler.setFormatter(formatter)
-      logger.addHandler(handler)
-      self.log_handler = handler
-
-    self.logger = logger
-    return logger
-
-  def cleanup_logger(self):
-    """Removes the job's logging handler."""
-    if self.logger and self.log_handler:
-      self.logger.removeHandler(self.log_handler)
-      self.log_handler = None
-      self.logger = None
+    def __init__(self, status: JobStatus):
+        self.id = str(f"ftjob-{uuid.uuid4()}")
+        self.status = status
+        self.fine_tuned_model = None
+        self.events: list[JobEvent] = []
+        self.checkpoints: list[JobCheckpoint] = []
+
+        self.created_at = datetime.now()
+
+    def add_event(self, event: JobEvent):
+        self.events.append(event)
+
+    def get_events(self) -> list[JobEvent]:
+        return self.events
+
+    def add_checkpoint(self, checkpoint: JobCheckpoint):
+        self.checkpoints.append(checkpoint)
+
+    def get_checkpoints(self) -> list[JobCheckpoint]:
+        return self.checkpoints
+
 
 class JobManager:
-  def __init__(self):
-    self.jobs = {}
-
-  def get_job(self, job_id: str):
-    if job_id not in self.jobs:
-      raise ValueError(f"Job {job_id} not found")
-    return self.jobs[job_id]
-
-  def create_job(self):
-    job = Job(id=str(uuid.uuid4()), status=JobStatus.PENDING)
-    self.jobs[job.id] = job
-    return job
+    def __init__(self, settings: Settings):
+        self.jobs = {}
+
+    def get_job(self, job_id: str):
+        if job_id not in self.jobs:
+            raise ValueError(f"Job {job_id} not found")
+        return self.jobs[job_id]
+
+    def create_job(self):
+        job = Job(status=JobStatus.PENDING)
+        self.jobs[job.id] = job
+        return job
+
+    def get_jobs(self):
+        return list(self.jobs.values())
+
+    def get_active_job(self):
+        for job in self.jobs.values():
+            if job.status == JobStatus.RUNNING:
+                return job
+        return None