gen-worker 0.1.4__py3-none-any.whl

gen_worker/torch_manager/utils/parse_cli.py

@@ -0,0 +1,98 @@
+import argparse
+import json
+import os
+import logging
+from .base_types.config import RuntimeConfig, PipelineConfig
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+# ====== Parse cli arguments ======
+
+
+def parse_pipeline_defs(value: Optional[str]) -> dict[str, PipelineConfig]:
+    """Parse pipeline definitions from command line argument"""
+    if not value:
+        return {}
+
+    try:
+        loaded = json.loads(value)
+        if isinstance(loaded, dict):
+            return loaded
+        else:
+            logger.error("Pipeline definitions are not a dictionary")
+            return {}
+    except json.JSONDecodeError as e:
+        logger.error(f"Failed to parse pipeline definitions: {e}")
+        return {}
+
+
+def parse_enabled_models(value: Optional[str]) -> list[str]:
+    """Parse enabled models from command line argument"""
+    if not value:
+        return []
+
+    try:
+        return json.loads(value)
+    except json.JSONDecodeError:
+        return value.split(",")
+
+
+def parse_arguments() -> RuntimeConfig:
+    """Parse command line arguments and return configuration"""
+    parser = argparse.ArgumentParser(description="Cozy Creator")
+
+    default_config = RuntimeConfig()
+
+    parser.add_argument(
+        "--home-dir",
+        default=default_config.home_dir,
+        help="Cozy creator's home directory",
+    )
+    parser.add_argument(
+        "--environment",
+        default=default_config.environment,
+        help="Server environment (dev/prod)",
+    )
+    parser.add_argument(
+        "--host", default=default_config.host, help="Hostname or IP-address"
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=default_config.port,
+        help="Port to bind Python runtime to",
+    )
+    parser.add_argument(
+        "--pipeline-defs",
+        type=str,
+        default=default_config.pipeline_defs,
+        help="JSON string of pipeline definitions",
+    )
+    parser.add_argument(
+        "--enabled-models",
+        type=str,
+        default=default_config.enabled_models,
+        help="Comma-separated list or JSON array of models to warm up",
+    )
+    parser.add_argument(
+        "--models-path",
+        type=str,
+        default=default_config.models_path,
+        help="Path to models directory",
+    )
+
+    args = parser.parse_args()
+
+    # Update config with parsed arguments
+    config = RuntimeConfig(
+        home_dir=args.home_dir,
+        environment=args.environment,
+        host=args.host,
+        port=args.port,
+        pipeline_defs=parse_pipeline_defs(args.pipeline_defs),
+        enabled_models=parse_enabled_models(args.enabled_models),
+        models_path=args.models_path or os.path.join(args.home_dir, "models"),
+    )
+
+    return config

gen_worker/torch_manager/utils/paths.py

@@ -0,0 +1,22 @@
+import os
+from .config import get_config
+
+
+
+def get_assets_dir():
+    config = get_config()
+    if config.assets_path:
+        return os.path.expanduser(config.assets_path)
+    return os.path.join(get_home_dir(), "assets")
+
+
+def get_models_dir():
+    config = get_config()
+
+    if config.models_path:
+        return os.path.expanduser(config.models_path)
+    return os.path.join(get_home_dir(), "models")
+
+
+def get_home_dir():
+    return os.path.expanduser(get_config().home_dir)

gen_worker/torch_manager/utils/repository.py

@@ -0,0 +1,141 @@
+import json
+import logging
+from typing import List, Dict, Any, Optional
+from decimal import Decimal
+
+logger = logging.getLogger(__name__)
+
+class PipelineDef:
+    """
+    Model class for pipeline definitions from the database.
+    """
+    def __init__(self, name: str, source: str = "", class_name: str = "", 
+                 custom_pipeline: str = "", default_args: Dict = None, 
+                 metadata: Dict = None, components: Dict = None,
+                 prompt_def: Optional['PromptDef'] = None, estimated_size_gb: Optional[float] = None):
+        self.name = name
+        self.source = source
+        self.class_name = class_name
+        self.custom_pipeline = custom_pipeline
+        self.default_args = default_args or {}
+        self.metadata = metadata or {}
+        self.components = components or {}
+        self.prompt_def = prompt_def
+        self.estimated_size_gb = estimated_size_gb
+
+class PromptDef:
+    """
+    Model class for prompt definitions from the database.
+    """
+    def __init__(self, pipeline_id: int, positive_prompt: str = "", negative_prompt: str = ""):
+        self.pipeline_id = pipeline_id
+        self.positive_prompt = positive_prompt
+        self.negative_prompt = negative_prompt
+
+def get_pipeline_defs(db_conn, pipeline_names: List[str]) -> List[PipelineDef]:
+    """
+    Retrieves pipeline definitions from the database based on their names.
+    Similar to the Go GetPipelineDefs function.
+    
+    Args:
+        db_conn: Database connection
+        pipeline_names: List of pipeline names to retrieve
+        
+    Returns:
+        List of PipelineDef objects
+    """
+    if not pipeline_names:
+        return []
+    
+    try:
+        pipeline_defs = []
+        with db_conn.cursor() as cur:
+            # Query to get pipeline definitions with their prompt definitions
+            query = """
+                SELECT 
+                    p.id,
+                    p.name,
+                    p.source,
+                    p.class_name,
+                    p.custom_pipeline,
+                    p.default_args,
+                    p.metadata,
+                    p.components,
+                    p.estimated_size_bytes,
+                    pr.positive_prompt,
+                    pr.negative_prompt
+                FROM 
+                    pipeline_defs p
+                LEFT JOIN 
+                    prompt_defs pr ON p.prompt_def_id = pr.id
+                WHERE 
+                    p.name = ANY(%s)
+            """
+            
+            cur.execute(query, (pipeline_names,))
+            rows = cur.fetchall()
+            
+            for row in rows:
+                # Create prompt_def if available
+                prompt_def = None
+                if row['positive_prompt'] or row['negative_prompt']:
+                    prompt_def = PromptDef(
+                        pipeline_id=row['id'],
+                        positive_prompt=row['positive_prompt'] or "",
+                        negative_prompt=row['negative_prompt'] or ""
+                    )
+                
+                # Parse JSON fields
+                default_args = {}
+                if row['default_args']:
+                    try:
+                        default_args = json.loads(row['default_args']) if isinstance(row['default_args'], str) else row['default_args']
+                    except json.JSONDecodeError:
+                        logger.warning(f"Failed to parse default_args for pipeline {row['name']}")
+                
+                metadata = {}
+                if row['metadata']:
+                    try:
+                        metadata = json.loads(row['metadata']) if isinstance(row['metadata'], str) else row['metadata']
+                    except json.JSONDecodeError:
+                        logger.warning(f"Failed to parse metadata for pipeline {row['name']}")
+                
+                components = {}
+                if row['components']:
+                    try:
+                        components = json.loads(row['components']) if isinstance(row['components'], str) else row['components']
+                    except json.JSONDecodeError:
+                        logger.warning(f"Failed to parse components for pipeline {row['name']}")
+
+                estimated_size_val_gb = None # Changed name
+                if row['estimated_size_bytes'] is not None: # DB column name is still estimated_size_bytes
+                    try:
+                        if isinstance(row['estimated_size_bytes'], Decimal):
+                            estimated_size_val_gb = float(row['estimated_size_bytes'])
+                        else:
+                            estimated_size_val_gb = float(str(row['estimated_size_bytes']))
+                    except (ValueError, TypeError) as e:
+                        logger.warning(f"Could not convert estimated_size_bytes ('{row['estimated_size_bytes']}') to float for {row['name']}: {e}")
+                
+                # Create PipelineDef
+                pipeline_def = PipelineDef(
+                    name=row['name'],
+                    source=row['source'] or "",
+                    class_name=row['class_name'] or "",
+                    custom_pipeline=row['custom_pipeline'] or "",
+                    default_args=default_args,
+                    metadata=metadata,
+                    components=components,
+                    prompt_def=prompt_def,
+                    estimated_size_gb=estimated_size_val_gb
+                )
+                
+                pipeline_defs.append(pipeline_def)
+        
+        return pipeline_defs
+    
+    except Exception as e:
+        logger.error(f"Error retrieving pipeline definitions: {e}")
+        # Ensure transaction is rolled back
+        db_conn.rollback()
+        raise

gen_worker/torch_manager/utils/utils.py

@@ -0,0 +1,43 @@
+import re
+from typing import Dict, Any
+from dataclasses import asdict, is_dataclass
+from .base_types.config import RuntimeConfig
+
+
+def flatten_architectures(architectures):
+    flat_architectures = {}
+    for arch_id, architecture in architectures.items():
+        if isinstance(architecture, list):
+            for arch in architecture:
+                flat_architectures[f"{arch_id}:{arch.__name__}"] = architecture
+        else:
+            flat_architectures[arch_id] = architecture
+
+    return flat_architectures
+
+
+def to_snake_case(value):
+    """
+    Convert CamelCase to snake_case
+    """
+    pattern = re.compile(r"(?<=[a-z])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])")
+    return pattern.sub("_", value).lower()
+
+
+def serialize_config(config: RuntimeConfig) -> dict[str, Any]:
+    """
+    Serialize a dataclass (like RuntimeConfig) into a dictionary.
+    This function handles nested dataclasses and converts them to dictionaries as well.
+    """
+
+    def serialize(obj):
+        if is_dataclass(obj):
+            return {k: serialize(v) for k, v in asdict(obj).items()}
+        elif isinstance(obj, list):
+            return [serialize(item) for item in obj]
+        elif isinstance(obj, dict):
+            return {k: serialize(v) for k, v in obj.items()}
+        else:
+            return obj
+
+    return serialize(config)

gen_worker/types.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+import os
+from typing import IO, Optional
+
+import msgspec
+
+
+class Asset(msgspec.Struct):
+    """Reference to a file in the invoking tenant's file store.
+
+    The worker runtime should populate `local_path` before invoking tenant code
+    so the function can open/read the file efficiently.
+    """
+
+    ref: str
+    tenant_id: Optional[str] = None
+    local_path: Optional[str] = None
+    mime_type: Optional[str] = None
+    size_bytes: Optional[int] = None
+    sha256: Optional[str] = None
+
+    def __fspath__(self) -> str:
+        if self.local_path is None:
+            raise ValueError("Asset.local_path is not set (file not materialized)")
+        return self.local_path
+
+    def open(self, mode: str = "rb") -> IO[bytes]:
+        if "b" not in mode:
+            raise ValueError("Asset.open only supports binary modes")
+        if self.local_path is None:
+            raise ValueError("Asset.local_path is not set (file not materialized)")
+        return open(self.local_path, mode)
+
+    def exists(self) -> bool:
+        if self.local_path is None:
+            return False
+        return os.path.exists(self.local_path)
+
+    def read_bytes(self, max_bytes: Optional[int] = None) -> bytes:
+        if self.local_path is None:
+            raise ValueError("Asset.local_path is not set (file not materialized)")
+        with open(self.local_path, "rb") as f:
+            data = f.read() if max_bytes is None else f.read(max_bytes + 1)
+        if max_bytes is not None and len(data) > max_bytes:
+            raise ValueError("asset too large to read into memory")
+        return data