Add ability to wait for an environment

.github/workflows/ci.yaml

@@ -118,7 +118,7 @@ jobs:
     name: "Test integration"
     needs: build-python
     runs-on: ubuntu-latest-16-cores
-    timeout-minutes: 10
+    timeout-minutes: 20
     steps:
       - uses: actions/checkout@v4
         with:

.gitignore

@@ -18,3 +18,4 @@ docs/README.md
 docs/CONTRIBUTING.md
 venv
 base-image
+flag_file

pkg/config/torch_compatibility_matrix.json

@@ -191,7 +191,8 @@
       "3.8",
       "3.9",
       "3.10",
-      "3.11"
+      "3.11",
+      "3.12"
     ]
   },
   {

python/cog/__init__.py

@@ -1,6 +1,6 @@
 from pydantic import BaseModel
 
-from .predictor import BasePredictor
+from .base_predictor import BasePredictor
 from .types import ConcatenateIterator, File, Input, Path, Secret
 
 try:

python/cog/base_input.py

@@ -0,0 +1,32 @@
+from pathlib import Path
+
+from pydantic import BaseModel
+
+from .types import (
+    URLPath,
+)
+
+
+# Base class for inputs, constructed dynamically in get_input_type().
+# (This can't be a docstring or it gets passed through to the schema.)
+class BaseInput(BaseModel):
+    class Config:
+        # When using `choices`, the type is converted into an enum to validate
+        # But, after validation, we want to pass the actual value to predict(), not the enum object
+        use_enum_values = True
+
+    def cleanup(self) -> None:
+        """
+        Cleanup any temporary files created by the input.
+        """
+        for _, value in self:
+            # Handle URLPath objects specially for cleanup.
+            # Also handle pathlib.Path objects, which cog.Path is a subclass of.
+            # A pathlib.Path object shouldn't make its way here,
+            # but both have an unlink() method, so we may as well be safe.
+            if isinstance(value, (URLPath, Path)):
+                # TODO: use unlink(missing_ok=...) when we drop Python 3.7 support.
+                try:
+                    value.unlink()
+                except FileNotFoundError:
+                    pass

python/cog/base_predictor.py

@@ -0,0 +1,26 @@
+from abc import ABC, abstractmethod
+from typing import Any, Optional, Union
+
+from .types import (
+    File as CogFile,
+)
+from .types import (
+    Path as CogPath,
+)
+
+
+class BasePredictor(ABC):
+    def setup(
+        self,
+        weights: Optional[Union[CogFile, CogPath, str]] = None,  # pylint: disable=unused-argument
+    ) -> None:
+        """
+        An optional method to prepare the model so multiple predictions run efficiently.
+        """
+        return
+
+    @abstractmethod
+    def predict(self, **kwargs: Any) -> Any:
+        """
+        Run a single prediction on the model
+        """

python/cog/code_xforms.py

@@ -1,7 +1,7 @@
 import ast
 import re
 import types
-from typing import Optional, Set, Union
+from typing import List, Optional, Set, Tuple, Union
 
 COG_IMPORT_MODULES = {"cog", "typing", "sys", "os", "functools", "pydantic", "numpy"}
 
@@ -67,7 +67,9 @@ def visit_FunctionDef(self, node: ast.FunctionDef) -> None:  # pylint: disable=i
     return extractor.function_source if extractor.function_source else ""
 
 
-def make_class_methods_empty(source_code: Union[str, ast.AST], class_name: str) -> str:
+def make_class_methods_empty(
+    source_code: Union[str, ast.AST], class_name: str, globals: List[ast.Assign]
+) -> Tuple[str, List[ast.Assign]]:
     """
     Transforms the source code of a specified class to remove the bodies of all its methods
     and replace them with 'return None'.
@@ -79,21 +81,42 @@ def make_class_methods_empty(source_code: Union[str, ast.AST], class_name: str)
     """
 
     class MethodBodyTransformer(ast.NodeTransformer):
+        def __init__(self, globals: List[ast.Assign]) -> None:
+            self.used_globals = set()
+            self._targets = {
+                target.id: global_name
+                for global_name in globals
+                for target in global_name.targets
+                if isinstance(target, ast.Name)
+            }
+
         def visit_ClassDef(self, node: ast.ClassDef) -> Optional[ast.AST]:  # pylint: disable=invalid-name
             if node.name == class_name:
                 for body_item in node.body:
                     if isinstance(body_item, ast.FunctionDef):
                         # Replace the body of the method with `return None`
                         body_item.body = [ast.Return(value=ast.Constant(value=None))]
+                        # Remove decorators from the function
+                        body_item.decorator_list = []
+                        # Determine if one our globals is referenced by the function.
+                        for default in body_item.args.defaults:
+                            if isinstance(default, ast.Call):
+                                for keyword in default.keywords:
+                                    if isinstance(keyword.value, ast.Name):
+                                        corresponding_global = self._targets.get(
+                                            keyword.value.id
+                                        )
+                                        if corresponding_global is not None:
+                                            self.used_globals.add(corresponding_global)
                 return node
 
             return None
 
     tree = source_code if isinstance(source_code, ast.AST) else ast.parse(source_code)
-    transformer = MethodBodyTransformer()
+    transformer = MethodBodyTransformer(globals)
     transformed_tree = transformer.visit(tree)
     class_code = ast.unparse(transformed_tree)
-    return class_code
+    return class_code, list(transformer.used_globals)
 
 
 def extract_method_return_type(
@@ -215,6 +238,17 @@ def visit_ImportFrom(self, node: ast.ImportFrom) -> None:  # pylint: disable=inv
     return "\n".join(extractor.imports)
 
 
+def _extract_globals(source_code: Union[str, ast.AST]) -> List[ast.Assign]:
+    tree = source_code if isinstance(source_code, ast.AST) else ast.parse(source_code)
+    if isinstance(tree, ast.Module):
+        return [x for x in tree.body if isinstance(x, ast.Assign)]
+    return []
+
+
+def _render_globals(globals: List[ast.Assign]) -> str:
+    return "\n".join([ast.unparse(x) for x in globals])
+
+
 def strip_model_source_code(
     source_code: str, class_name: str, method_name: str
 ) -> Optional[str]:
@@ -234,14 +268,22 @@ def strip_model_source_code(
     class_source = (
         None if not class_name else extract_class_source(source_code, class_name)
     )
+    globals = _extract_globals(source_code)
     if class_source:
-        class_source = make_class_methods_empty(class_source, class_name)
+        class_source, globals = make_class_methods_empty(
+            class_source, class_name, globals
+        )
         return_type = extract_method_return_type(class_source, class_name, method_name)
         return_class_source = (
             extract_class_source(source_code, return_type) if return_type else ""
         )
-        model_source = (
-            imports + "\n\n" + return_class_source + "\n\n" + class_source + "\n"
+        rendered_globals = _render_globals(globals)
+        model_source = "\n".join(
+            [
+                x
+                for x in [imports, rendered_globals, return_class_source, class_source]
+                if x
+            ]
         )
     else:
         # use class_name specified in cog.yaml as method_name
@@ -256,7 +298,5 @@ def strip_model_source_code(
         return_class_source = (
             extract_class_source(source_code, return_type) if return_type else ""
         )
-        model_source = (
-            imports + "\n\n" + return_class_source + "\n\n" + function_source + "\n"
-        )
+        model_source = "\n".join([imports, return_class_source, function_source])
     return model_source

python/cog/command/openapi_schema.py

@@ -7,8 +7,8 @@
 import json
 from typing import Any, Dict, List, Union
 
+from ..config import Config
 from ..errors import CogError, ConfigDoesNotExist, PredictorNotSet
-from ..predictor import load_config
 from ..schema import Status
 from ..server.http import create_app
 from ..suppress_output import suppress_output
@@ -37,8 +37,7 @@ def remove_title_next_to_ref(
     schema = {}
     try:
         with suppress_output():
-            config = load_config()
-            app = create_app(config, shutdown_event=None, is_build=True)
+            app = create_app(cog_config=Config(), shutdown_event=None, is_build=True)
             if (
                 app.state.setup_result
                 and app.state.setup_result.status == Status.FAILED

python/cog/config.py

@@ -0,0 +1,144 @@
+import os
+import sys
+import uuid
+from typing import Optional, Tuple, Type
+
+import structlog
+import yaml
+from pydantic import BaseModel
+
+from .base_input import BaseInput
+from .base_predictor import BasePredictor
+from .code_xforms import load_module_from_string, strip_model_source_code
+from .env_property import env_property
+from .errors import ConfigDoesNotExist
+from .mode import Mode
+from .predictor import (
+    get_input_type,
+    get_output_type,
+    get_predictor,
+    get_training_input_type,
+    get_training_output_type,
+    load_full_predictor_from_file,
+)
+from .types import CogConfig
+from .wait import wait_for_env
+
+COG_YAML_FILE = "cog.yaml"
+COG_PREDICTOR_PREDICT_ENV_VAR = "COG_PREDICTOR_PREDICT"
+COG_PREDICTOR_TRAIN_ENV_VAR = "COG_PREDICTOR_TRAIN"
+COG_PREDICTOR_ENV_VAR = "COG_PREDICTOR"
+COG_GPU_ENV_VAR = "COG_GPU"
+PREDICT_METHOD_NAME = "predict"
+TRAIN_METHOD_NAME = "train"
+
+log = structlog.get_logger("cog.config")
+
+
+def _method_name_from_mode(mode: Mode) -> str:
+    if mode == Mode.PREDICT:
+        return PREDICT_METHOD_NAME
+    elif mode == Mode.TRAIN:
+        return TRAIN_METHOD_NAME
+    raise ValueError(f"Mode {mode} not recognised for method name mapping")
+
+
+class Config:
+    """A class for reading the cog.yaml properties."""
+
+    def __init__(self, config: Optional[CogConfig] = None) -> None:
+        self._config = config
+
+    @property
+    def _cog_config(self) -> CogConfig:
+        """
+        Warning: Do not access this directly outside this class, instead
+        write an explicit public property and back it by an @env_property
+        to allow for the possibility of injecting the property you are
+        trying to read without relying on the underlying file.
+        """
+        config = self._config
+        if config is None:
+            wait_for_env(include_imports=False)
+            config_path = os.path.abspath(COG_YAML_FILE)
+            try:
+                with open(config_path, encoding="utf-8") as handle:
+                    config = yaml.safe_load(handle)
+            except FileNotFoundError as e:
+                raise ConfigDoesNotExist(
+                    f"Could not find {config_path}",
+                ) from e
+            self._config = config
+        return config
+
+    @property
+    @env_property(COG_PREDICTOR_PREDICT_ENV_VAR)
+    def predictor_predict_ref(self) -> Optional[str]:
+        """Find the predictor ref for the predict mode."""
+        return self._cog_config.get(str(Mode.PREDICT))
+
+    @property
+    @env_property(COG_PREDICTOR_TRAIN_ENV_VAR)
+    def predictor_train_ref(self) -> Optional[str]:
+        """Find the predictor ref for the train mode."""
+        return self._cog_config.get(str(Mode.TRAIN))
+
+    @property
+    @env_property(COG_GPU_ENV_VAR)
+    def requires_gpu(self) -> bool:
+        """Whether this cog requires the use of a GPU."""
+        return bool(self._cog_config.get("build", {}).get("gpu", False))
+
+    def _predictor_code(
+        self, module_path: str, class_name: str, method_name: str
+    ) -> Optional[str]:
+        source_code = os.environ.get(COG_PREDICTOR_ENV_VAR)
+        if source_code is not None:
+            return source_code
+        if sys.version_info >= (3, 9):
+            wait_for_env(include_imports=False)
+            with open(module_path, encoding="utf-8") as file:
+                return strip_model_source_code(file.read(), class_name, method_name)
+        return None
+
+    def _load_predictor_for_types(self, ref: str, method_name: str) -> BasePredictor:
+        module_path, class_name = ref.split(":", 1)
+        code = self._predictor_code(module_path, class_name, method_name)
+        module_name = os.path.basename(module_path).split(".py", 1)[0]
+        module = None
+        if code is not None:
+            try:
+                module = load_module_from_string(uuid.uuid4().hex, code)
+            except Exception as e:  # pylint: disable=broad-exception-caught
+                log.info(f"[{module_name}] fast loader failed: {e}")
+        if module is None:
+            wait_for_env(include_imports=False)
+            module = load_full_predictor_from_file(module_path, module_name)
+        return get_predictor(module, class_name)
+
+    def get_predictor_ref(self, mode: Mode) -> str:
+        """Find the predictor reference for a given mode."""
+        predictor_ref = None
+        if mode == Mode.PREDICT:
+            predictor_ref = self.predictor_predict_ref
+        elif mode == Mode.TRAIN:
+            predictor_ref = self.predictor_train_ref
+        if predictor_ref is None:
+            raise ValueError(f"Could not find predictor ref for mode {mode}")
+        return predictor_ref
+
+    def get_predictor_types(
+        self, mode: Mode
+    ) -> Tuple[Type[BaseInput], Type[BaseModel]]:
+        """Find the input and output types of a predictor."""
+        predictor_ref = self.get_predictor_ref(mode=mode)
+        predictor = self._load_predictor_for_types(
+            predictor_ref, _method_name_from_mode(mode=mode)
+        )
+        if mode == Mode.PREDICT:
+            return get_input_type(predictor), get_output_type(predictor)
+        elif mode == Mode.TRAIN:
+            return get_training_input_type(predictor), get_training_output_type(
+                predictor
+            )
+        raise ValueError(f"Mode {mode} not found for generating input/output types.")