Add ability to wait for an environment (#1957)

* Send in app threads directly from args

* This allows us not to read the config for the
app threads until absolutely necessary.

* Load config right before necessary

* Delay loading the cog yaml until it is read from
in the code.

* Add waiting for a wait file

* Allow environment variables to control whether
we wait for a file to appear before further
processing
* This allows the python interpreter to boot up
while we wait for other files to become available

* Add wait_for_imports ability

* Allow waiting for a general environment by
importing select python packages while the system
boots up around cog.

* Fix lint on src_path

* Fix watchdog version

* Remove load_config in openapi_schema cmd

* Do not access root files on GHA workers

* Set recursive to true

* I shouldn’t need to do this but want to see if
it relieves the errors on linux.

* Watch the directory instead

* Add code_xforms test

* A small test to make sense of what the code
stripper is doing.

* Add http server to test to let it respond

* Abstract away cog config

* Create a class for accessing cog config
* Only access variables from the config by
properties
* Gate those properties with environment variable
function decorators to allow fetching the config
from the environment rather than a file.
* This allows the environment to begin running
without needing the /src

* Wait for environment before executing setup

* While waiting for the environment to boot, load
The designated imports to speed up interpreter
Time.

* Fix Type on lower python versions

* Skip test_strip_model_source_code if < 3.9

 * In strip model source code we use an AST
function that isn’t available prior to 3.9

* Change COG_WAIT_IMPORTS to COG_EAGER_IMPORTS

* We aren’t waiting for these imports, we are
loading them eagerly while we wait.

* Bump integration test timeout to 20 mins

* We have so many that we need to increase this.

* Add tests for Config class

* Fix an issue with env_property where it could
not handle Optional or Union

* Fix get_args and get_origin in python 3.7

* Create these functions because they don’t exist
in python 3.7

* Add more tests for config

* Check wait flag has fallen before eager import

* Add watch handler tests

* Remove watchdog and use SIGUSR2 for signalling

* Fix no torch import in tests

* Do naive waiting for file

* Do not wait for a signal, just use a while loop
with a 10 ms interval to check for presence of
file.

* Consolidate code_xforms tests

* test_code_xforms actually lived in the server
directory and not in the cog directory, move these
tests into the cog directory.
* Add a test and code for removing function
decorators from the target functions, these can
remove the ability to load the slimmed predictor.

* Add logic to keep referenced globals

* Keep a global if it is referenced as a default
from a functions arguments.

* Convert set of globals to list

* Add logging for file waiting

* Fix lint issue by checking that tree is ast.Module

* Use typing.List instead of list for older python

* Fix more List issues

* Add further logging if the module fast loader failed

* Change tuple to Tuple for older pythons

* Add setup logs to mirror prediction logs

* Allow the logs to show when certain setup steps
are happening

* Add 3.12 to torch 2.3.0

* Add insertion of python path

* Add support for handling subclasses in slim predict

* Fix listing issues

* Wrap in str

* Remove extraneous print

* Fix python path race condition

* Remove predictor import

* Fix pydantic 2 errors

* More pydantic 2 fixes

* Support multiple class names and method names

* This allows consumers to strip source code but
Maintain multiple classes and methods
* This is important if they want to maintain both
Predict and train types

* Use lists instead single instances for stripping

* Fix types imports

* Revert "Add 3.12 to torch 2.3.0"

This reverts commit 096c45b.

* Fix imports after merge

* Add COG_TRAIN_PREDICTOR

* Ensure that train/predictor stripped code is
kept separate in the event a user uses separate
files to refer to these different modes.

* Rename env vars to be sensical

* Old ones used to be many overloaded concepts.

* Add PYTHONPATH environment variables

* We can’t just update the system path, we also
need to change the PYTHONPATH environment variable
* This allows forked processes to see this
environment variable and act accordingly in user
code.

* Add consistent debug logging in config

* Some tests previously checked the debug logging
* Make sure we conform to the same debug logging

* Add test_strip_model_source_code_keeps_referenced_class_from_function

* Confirm that this behaviour is consistent with
functions.

* Explicitly check return code in test train

* Allows better debugging of stderr et al

* Handle response types in _predict

* Currently predict assumes that PredictionResponse
is the serialisation target.
* This isn’t the case if we are calling _predict
from a training endpoint
* Allow response_type to be fed into the _predict
method to inform it of what kind of response it
should expect.

* Make cog train a first class CLI function

* Use the proper endpoints to call training
functions.
* Log the command properly to the user.

* Add back missing imports from merge

* Remove connection import

* Fix lint

* Add R8_ prefix to PYTHON_VERSION (#2058)

---------

Signed-off-by: Will Sackfield <[email protected]>
Co-authored-by: Neville Li <[email protected]>

Author: 8W9aG

.gitignore

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

python/cog/config.py

@@ -22,6 +22,7 @@
     load_full_predictor_from_file,
 )
 from .types import CogConfig
+from .wait import wait_for_env
 
 COG_YAML_FILE = "cog.yaml"
 COG_PREDICT_TYPE_STUB_ENV_VAR = "COG_PREDICT_TYPE_STUB"
@@ -67,6 +68,7 @@ def _cog_config(self) -> CogConfig:
         """
         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:
@@ -108,6 +110,7 @@ def _predictor_code(
         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])
         else:
@@ -130,6 +133,7 @@ def _load_predictor_for_types(
                 log.info(f"[{module_name}] fast loader failed: {e}")
         if module is None:
             log.debug(f"[{module_name}] falling back to slow loader")
+            wait_for_env(include_imports=False)
             module = load_full_predictor_from_file(module_path, module_name)
         return get_predictor(module, class_name)
 

python/cog/server/worker.py

@@ -30,6 +30,7 @@
 from ..json import make_encodeable
 from ..predictor import get_predict, load_predictor_from_ref, run_setup
 from ..types import PYDANTIC_V2, URLPath
+from ..wait import wait_for_env
 from .connection import AsyncConnection, LockedConnection
 from .eventtypes import (
     Cancel,
@@ -371,6 +372,7 @@ def _current_tag(self) -> Optional[str]:
 
     def _setup(self, redirector: AsyncStreamRedirector) -> None:
         done = Done()
+        wait_for_env()
         try:
             self._predictor = load_predictor_from_ref(self._predictor_ref)
             # Could be a function or a class

python/cog/wait.py

@@ -0,0 +1,81 @@
+import importlib
+import os
+import sys
+import time
+
+import structlog
+
+COG_WAIT_FILE_ENV_VAR = "COG_WAIT_FILE"
+COG_EAGER_IMPORTS_ENV_VAR = "COG_EAGER_IMPORTS"
+COG_PYENV_PATH_ENV_VAR = "COG_PYENV_PATH"
+PYTHONPATH_ENV_VAR = "PYTHONPATH"
+PYTHON_VERSION_ENV_VAR = "R8_PYTHON_VERSION"
+
+log = structlog.get_logger("cog.wait")
+
+
+def _wait_flag_fallen() -> bool:
+    wait_file = os.environ.get(COG_WAIT_FILE_ENV_VAR)
+    if wait_file is None:
+        return True
+    return os.path.exists(wait_file)
+
+
+def _insert_pythonpath() -> None:
+    pyenv_path = os.environ.get(COG_PYENV_PATH_ENV_VAR)
+    if pyenv_path is None:
+        return
+    full_module_path = os.path.join(
+        pyenv_path,
+        "lib",
+        "python" + os.environ[PYTHON_VERSION_ENV_VAR],
+        "site-packages",
+    )
+    if full_module_path not in sys.path:
+        sys.path.append(full_module_path)
+    os.environ[PYTHONPATH_ENV_VAR] = ":".join(sys.path)
+
+
+def wait_for_file(timeout: float = 60.0) -> bool:
+    """Wait for a file in the environment variables."""
+    wait_file = os.environ.get(COG_WAIT_FILE_ENV_VAR)
+    if wait_file is None:
+        return True
+    if os.path.exists(wait_file):
+        log.info(f"Wait file found {wait_file}...")
+        return True
+    log.info(f"Waiting for file {wait_file}...")
+    time_taken = 0.0
+    while time_taken < timeout:
+        sleep_time = 0.01
+        time.sleep(sleep_time)
+        time_taken += sleep_time
+        if os.path.exists(wait_file):
+            return True
+    log.info(f"Waiting for file {wait_file} timed out.")
+    return False
+
+
+def eagerly_import_modules() -> int:
+    """Wait for python to import big modules."""
+    wait_imports = os.environ.get(COG_EAGER_IMPORTS_ENV_VAR)
+    import_count = 0
+    if wait_imports is None:
+        return import_count
+    log.info(f"Eagerly importing {wait_imports}.")
+    for import_statement in wait_imports.split(","):
+        importlib.import_module(import_statement)
+        import_count += 1
+    return import_count
+
+
+def wait_for_env(file_timeout: float = 60.0, include_imports: bool = True) -> bool:
+    """Wait for the environment to load."""
+    if _wait_flag_fallen():
+        _insert_pythonpath()
+        return True
+    if include_imports:
+        eagerly_import_modules()
+    waited = wait_for_file(timeout=file_timeout)
+    _insert_pythonpath()
+    return waited

python/tests/test_config.py

@@ -130,10 +130,12 @@ def test_get_predictor_types_with_env_var():
 from cog import BasePredictor, Path
 from typing import Optional
 from pydantic import BaseModel
+
 class ModelOutput(BaseModel):
     success: bool
     error: Optional[str]
     segmentedImage: Optional[Path]
+
 class Predictor(BasePredictor):
     def predict(self, msg: str) -> ModelOutput:
         return None

python/tests/test_wait.py

@@ -0,0 +1,123 @@
+import os
+import sys
+import tempfile
+import threading
+import time
+from pathlib import Path
+
+from cog.wait import (
+    COG_EAGER_IMPORTS_ENV_VAR,
+    COG_PYENV_PATH_ENV_VAR,
+    COG_WAIT_FILE_ENV_VAR,
+    PYTHON_VERSION_ENV_VAR,
+    PYTHONPATH_ENV_VAR,
+    eagerly_import_modules,
+    wait_for_env,
+    wait_for_file,
+)
+
+
+def test_wait_for_file_no_env_var():
+    if COG_WAIT_FILE_ENV_VAR in os.environ:
+        del os.environ[COG_WAIT_FILE_ENV_VAR]
+    result = wait_for_file()
+    assert result, "We should immediately return when no wait file is specified."
+
+
+def test_wait_for_file_exists():
+    with tempfile.NamedTemporaryFile() as tmpfile:
+        os.environ[COG_WAIT_FILE_ENV_VAR] = tmpfile.name
+        result = wait_for_file(timeout=5.0)
+        del os.environ[COG_WAIT_FILE_ENV_VAR]
+        assert result, "We should immediately return when the file already exists."
+
+
+def test_wait_for_file_waits_for_file():
+    wait_file = os.path.join(os.path.dirname(__file__), "flag_file")
+    if os.path.exists(wait_file):
+        os.remove(wait_file)
+    os.environ[COG_WAIT_FILE_ENV_VAR] = wait_file
+
+    def create_file():
+        time.sleep(2.0)
+        Path(wait_file).touch()
+
+    thread = threading.Thread(target=create_file)
+    thread.start()
+    result = wait_for_file(timeout=5.0)
+    del os.environ[COG_WAIT_FILE_ENV_VAR]
+    os.remove(wait_file)
+    assert result, "We should return when the file is touched."
+
+
+def test_wait_for_file_timeout():
+    os.environ[COG_WAIT_FILE_ENV_VAR] = os.path.join(
+        os.path.dirname(__file__), "a_file_unknown"
+    )
+    result = wait_for_file(timeout=5.0)
+    del os.environ[COG_WAIT_FILE_ENV_VAR]
+    assert not result, "We should return false when the timeout triggers."
+
+
+def test_eagerly_import_modules_no_env_var():
+    if COG_EAGER_IMPORTS_ENV_VAR in os.environ:
+        del os.environ[COG_EAGER_IMPORTS_ENV_VAR]
+    eagerly_import_modules()
+
+
+def test_eagerly_import_modules():
+    os.environ[COG_EAGER_IMPORTS_ENV_VAR] = "pytest,pathlib,time"
+    import_count = eagerly_import_modules()
+    del os.environ[COG_EAGER_IMPORTS_ENV_VAR]
+    assert import_count == 3, "There should be 3 imports performed"
+
+
+def test_wait_for_env_no_env_vars():
+    if COG_WAIT_FILE_ENV_VAR in os.environ:
+        del os.environ[COG_WAIT_FILE_ENV_VAR]
+    if COG_EAGER_IMPORTS_ENV_VAR in os.environ:
+        del os.environ[COG_EAGER_IMPORTS_ENV_VAR]
+    result = wait_for_env()
+    assert (
+        result
+    ), "We should return true if we have no env vars associated with the wait."
+
+
+def test_wait_for_env():
+    with tempfile.NamedTemporaryFile() as tmpfile:
+        os.environ[COG_WAIT_FILE_ENV_VAR] = tmpfile.name
+        os.environ[COG_EAGER_IMPORTS_ENV_VAR] = "pytest,pathlib,time"
+        result = wait_for_env()
+        assert (
+            result
+        ), "We should return true if we have waited for the right environment."
+        del os.environ[COG_EAGER_IMPORTS_ENV_VAR]
+        del os.environ[COG_WAIT_FILE_ENV_VAR]
+
+
+def test_wait_inserts_pythonpath():
+    with tempfile.NamedTemporaryFile() as tmpfile:
+        original_sys_path = sys.path.copy()
+        original_python_path = os.environ.get(PYTHONPATH_ENV_VAR)
+        pyenv_path = os.path.dirname(tmpfile.name)
+        os.environ[COG_WAIT_FILE_ENV_VAR] = tmpfile.name
+        os.environ[COG_EAGER_IMPORTS_ENV_VAR] = "pytest,pathlib,time"
+        os.environ[COG_PYENV_PATH_ENV_VAR] = pyenv_path
+        os.environ[PYTHON_VERSION_ENV_VAR] = "3.11"
+        wait_for_env()
+        del os.environ[PYTHON_VERSION_ENV_VAR]
+        del os.environ[COG_PYENV_PATH_ENV_VAR]
+        del os.environ[COG_EAGER_IMPORTS_ENV_VAR]
+        del os.environ[COG_WAIT_FILE_ENV_VAR]
+        current_python_path = os.environ[PYTHONPATH_ENV_VAR]
+        if original_python_path is None:
+            del os.environ[PYTHONPATH_ENV_VAR]
+        else:
+            os.environ[PYTHONPATH_ENV_VAR] = original_python_path
+        sys.path = original_sys_path
+        expected_path = ":".join(
+            original_sys_path + [pyenv_path + "/lib/python3.11/site-packages"]
+        )
+        assert (
+            expected_path == current_python_path
+        ), "Our python path should be updated with the pyenv path."