From 610121aba79a7ef9901712db843371255e1ea400 Mon Sep 17 00:00:00 2001 From: Dan Buch Date: Sun, 17 Nov 2024 15:10:56 -0500 Subject: [PATCH] Rename to `coglet` plus reorg so that: - `python -m coglet` is the main entrypoint - `cog` compatibility imports are a thin wrapper around `coglet.api` and are only available for import once `coglet/_compat` is put into `sys.path`, thus greatly reducing the likelihood of import collisions. - `cog/internal/` directory collapsed up a level to `coglet/` since the above import collision problem is fixed. --- python/README.md | 11 ++- python/cog/internal/__init__.py | 0 python/coglet/__init__.py | 16 ++++ python/coglet/__main__.py | 80 +++++++++++++++++++ .../__init__.py => coglet/_compat/cog.py} | 9 ++- python/{cog/internal => coglet}/adt.py | 8 +- python/{cog => coglet}/api.py | 0 .../{cog/internal => coglet}/file_runner.py | 68 ++-------------- python/{cog/internal => coglet}/inspector.py | 11 ++- python/{cog/internal => coglet}/openapi.json | 0 python/{cog/internal => coglet}/runner.py | 7 +- python/{cog/internal => coglet}/schemas.py | 2 +- python/{cog/internal => coglet}/util.py | 7 +- python/pyproject.toml | 7 +- python/tests/test_file_runner.py | 4 +- python/tests/test_file_runner_async.py | 2 +- python/tests/test_file_runner_iterator.py | 2 +- python/tests/test_predictors.py | 10 +-- python/tests/test_util.py | 4 +- python/tests/test_weights.py | 2 +- 20 files changed, 152 insertions(+), 98 deletions(-) mode change 120000 => 100644 python/README.md delete mode 100644 python/cog/internal/__init__.py create mode 100644 python/coglet/__init__.py create mode 100644 python/coglet/__main__.py rename python/{cog/__init__.py => coglet/_compat/cog.py} (50%) rename python/{cog/internal => coglet}/adt.py (93%) rename python/{cog => coglet}/api.py (100%) rename python/{cog/internal => coglet}/file_runner.py (77%) rename python/{cog/internal => coglet}/inspector.py (96%) rename python/{cog/internal => coglet}/openapi.json (100%) rename python/{cog/internal => coglet}/runner.py (97%) rename python/{cog/internal => coglet}/schemas.py (99%) rename python/{cog/internal => coglet}/util.py (90%) diff --git a/python/README.md b/python/README.md deleted file mode 120000 index 32d46ee..0000000 --- a/python/README.md +++ /dev/null @@ -1 +0,0 @@ -../README.md \ No newline at end of file diff --git a/python/README.md b/python/README.md new file mode 100644 index 0000000..0e00058 --- /dev/null +++ b/python/README.md @@ -0,0 +1,10 @@ +# Coglet + +Coglet provides a minimum viable [Cog] runtime primarily for use within the Replicate +platform, e.g.: + +``` +python -m coglet --working-dir path/to/code/ --module-name predict --class-name Predictor +``` + +[Cog]: diff --git a/python/cog/internal/__init__.py b/python/cog/internal/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/python/coglet/__init__.py b/python/coglet/__init__.py new file mode 100644 index 0000000..213ade8 --- /dev/null +++ b/python/coglet/__init__.py @@ -0,0 +1,16 @@ +import pathlib +import sys +import warnings + +# NOTE: The compatibility import provided in `./_compat/cog.py` **SHOULD NOT** be in +# PYTHONPATH until `coglet` is imported. This prevents `coglet` from interfering with +# normal usage of `cog` within a given python environment. +warnings.warn( + ( + 'coglet/_compat/ is being added to the front of sys.path ' + "for 'cog' import compatibility" + ), + category=ImportWarning, + stacklevel=2, +) +sys.path.insert(0, str(pathlib.Path(__file__).absolute().parent / '_compat')) diff --git a/python/coglet/__main__.py b/python/coglet/__main__.py new file mode 100644 index 0000000..931dea9 --- /dev/null +++ b/python/coglet/__main__.py @@ -0,0 +1,80 @@ +import argparse +import asyncio +import contextvars +import logging +import sys +from typing import Optional + +from coglet import file_runner + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + '--working-dir', metavar='DIR', required=True, help='working directory' + ) + parser.add_argument( + '--module-name', metavar='NAME', required=True, help='Python module name' + ) + parser.add_argument( + '--class-name', metavar='NAME', required=True, help='Python class name' + ) + + _ctx_pid: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar( + 'pid', default=None + ) + _ctx_newline: contextvars.ContextVar[bool] = contextvars.ContextVar( + 'newline', default=False + ) + + logger = logging.getLogger('coglet') + logger.setLevel(logging.INFO) + handler = logging.StreamHandler(sys.stdout) + handler.setFormatter( + logging.Formatter( + '%(asctime)s\t%(levelname)s\t[%(name)s]\t%(filename)s:%(lineno)d\t%(message)s' + ) + ) + logger.addHandler(handler) + + _stdout_write = sys.stdout.write + _stderr_write = sys.stderr.write + + def _ctx_write(write_fn): + def _write(s: str) -> int: + pid = _ctx_pid.get() + if pid is None: + return write_fn(s) + else: + n = 0 + if _ctx_newline.get(): + n += write_fn(f'[pid={pid}] ') + if s[-1] == '\n': + _ctx_newline.set(True) + s = s[:-1].replace('\n', f'\n[pid={pid}] ') + '\n' + else: + _ctx_newline.set(False) + s = s.replace('\n', f'\n[pid={pid}] ') + n += write_fn(s) + return n + + return _write + + sys.stdout.write = _ctx_write(_stdout_write) # type: ignore + sys.stderr.write = _ctx_write(_stderr_write) # type: ignore + + args = parser.parse_args() + + return asyncio.run( + file_runner.FileRunner( + logger=logger, + working_dir=args.working_dir, + module_name=args.module_name, + class_name=args.class_name, + ctx_pid=_ctx_pid, + ).start() + ) + + +if __name__ == '__main__': + sys.exit(main()) diff --git a/python/cog/__init__.py b/python/coglet/_compat/cog.py similarity index 50% rename from python/cog/__init__.py rename to python/coglet/_compat/cog.py index 4a3da9c..44f30f8 100644 --- a/python/cog/__init__.py +++ b/python/coglet/_compat/cog.py @@ -1,4 +1,11 @@ -from cog.api import BaseModel, BasePredictor, ConcatenateIterator, Input, Path, Secret +from coglet.api import ( + BaseModel, + BasePredictor, + ConcatenateIterator, + Input, + Path, + Secret, +) __all__ = [ 'BaseModel', diff --git a/python/cog/internal/adt.py b/python/coglet/adt.py similarity index 93% rename from python/cog/internal/adt.py rename to python/coglet/adt.py index d424757..5b619c5 100644 --- a/python/cog/internal/adt.py +++ b/python/coglet/adt.py @@ -3,7 +3,7 @@ from enum import Enum from typing import Any, Dict, Iterator, List, Optional, Union -import cog +import coglet.api class Type(Enum): @@ -41,8 +41,8 @@ class Kind(Enum): float: Type.FLOAT, int: Type.INTEGER, str: Type.STRING, - cog.Path: Type.PATH, - cog.Secret: Type.SECRET, + coglet.api.Path: Type.PATH, + coglet.api.Secret: Type.SECRET, } # Cog types to JSON types @@ -68,7 +68,7 @@ class Kind(Enum): CONTAINER_TO_COG = { list: Kind.LIST, typing.get_origin(Iterator): Kind.ITERATOR, - cog.ConcatenateIterator: Kind.CONCAT_ITERATOR, + coglet.api.ConcatenateIterator: Kind.CONCAT_ITERATOR, } diff --git a/python/cog/api.py b/python/coglet/api.py similarity index 100% rename from python/cog/api.py rename to python/coglet/api.py diff --git a/python/cog/internal/file_runner.py b/python/coglet/file_runner.py similarity index 77% rename from python/cog/internal/file_runner.py rename to python/coglet/file_runner.py index e33757d..1a9225a 100644 --- a/python/cog/internal/file_runner.py +++ b/python/coglet/file_runner.py @@ -1,4 +1,3 @@ -import argparse import asyncio import contextvars import json @@ -9,7 +8,7 @@ import sys from typing import Any, Dict, Optional -from cog.internal import inspector, runner, schemas, util +from coglet import inspector, runner, schemas, util class FileRunner: @@ -26,16 +25,19 @@ class FileRunner: def __init__( self, + *, logger: logging.Logger, working_dir: str, module_name: str, class_name: str, + ctx_pid: contextvars.ContextVar[Optional[str]], ): self.logger = logger self.working_dir = working_dir self.module_name = module_name self.class_name = class_name self.runner: Optional[runner.Runner] = None + self.ctx_pid = ctx_pid self.isatty = sys.stdout.isatty() async def start(self) -> int: @@ -147,7 +149,7 @@ async def start(self) -> int: async def _predict(self, pid: str, req: Dict[str, Any]) -> None: assert self.runner is not None - _ctx_pid.set(pid) + self.ctx_pid.set(pid) resp: Dict[str, Any] = { 'started_at': util.now_iso(), 'status': 'starting', @@ -193,63 +195,3 @@ def _respond( def _signal(self, signum: int) -> None: if not self.isatty: os.kill(os.getppid(), signum) - - -parser = argparse.ArgumentParser() -parser.add_argument( - '--working-dir', metavar='DIR', required=True, help='working directory' -) -parser.add_argument( - '--module-name', metavar='NAME', required=True, help='Python module name' -) -parser.add_argument( - '--class-name', metavar='NAME', required=True, help='Python class name' -) - -_ctx_pid: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar( - 'pid', default=None -) -_ctx_newline: contextvars.ContextVar[bool] = contextvars.ContextVar( - 'newline', default=False -) - -if __name__ == '__main__': - logger = logging.getLogger('cog-file-runner') - logger.setLevel(logging.INFO) - handler = logging.StreamHandler(sys.stdout) - handler.setFormatter( - logging.Formatter( - '%(asctime)s\t%(levelname)s\t[%(name)s]\t%(filename)s:%(lineno)d\t%(message)s' - ) - ) - logger.addHandler(handler) - - _stdout_write = sys.stdout.write - _stderr_write = sys.stderr.write - - def _ctx_write(write_fn): - def _write(s: str) -> int: - pid = _ctx_pid.get() - if pid is None: - return write_fn(s) - else: - n = 0 - if _ctx_newline.get(): - n += write_fn(f'[pid={pid}] ') - if s[-1] == '\n': - _ctx_newline.set(True) - s = s[:-1].replace('\n', f'\n[pid={pid}] ') + '\n' - else: - _ctx_newline.set(False) - s = s.replace('\n', f'\n[pid={pid}] ') - n += write_fn(s) - return n - - return _write - - sys.stdout.write = _ctx_write(_stdout_write) # type: ignore - sys.stderr.write = _ctx_write(_stderr_write) # type: ignore - - args = parser.parse_args() - fr = FileRunner(logger, args.working_dir, args.module_name, args.class_name) - sys.exit(asyncio.run(fr.start())) diff --git a/python/cog/internal/inspector.py b/python/coglet/inspector.py similarity index 96% rename from python/cog/internal/inspector.py rename to python/coglet/inspector.py index 4630732..d4ef9c8 100644 --- a/python/cog/internal/inspector.py +++ b/python/coglet/inspector.py @@ -4,8 +4,7 @@ import typing from typing import Callable, Optional -import cog -from cog.internal import adt, util +from coglet import adt, api, util def _check_parent(child: type, parent: type) -> bool: @@ -38,7 +37,7 @@ def _validate_predict(f: Callable) -> None: def _validate_input( - name: str, cog_t: adt.Type, is_list: bool, cog_in: cog.Input + name: str, cog_t: adt.Type, is_list: bool, cog_in: api.Input ) -> None: defaults = [] if cog_in.default is not None: @@ -102,7 +101,7 @@ def _validate_input( def _input_adt( - order: int, name: str, tpe: type, cog_in: Optional[cog.Input] + order: int, name: str, tpe: type, cog_in: Optional[api.Input] ) -> adt.Input: cog_t, is_list = util.check_cog_type(tpe) assert cog_t is not None, f'unsupported input type for {name}' @@ -139,7 +138,7 @@ def _input_adt( def _output_adt(tpe: type) -> adt.Output: - if inspect.isclass(tpe) and _check_parent(tpe, cog.BaseModel): + if inspect.isclass(tpe) and _check_parent(tpe, api.BaseModel): assert tpe.__name__ == 'Output', 'output type must be named Output' fields = {} for name, t in tpe.__annotations__.items(): @@ -183,7 +182,7 @@ def create_predictor(module_name: str, class_name: str) -> adt.Predictor: cls = getattr(module, class_name) assert inspect.isclass(cls), f'not a class: {fullname}' assert _check_parent( - cls, cog.BasePredictor + cls, api.BasePredictor ), f'predictor {fullname} does not inherit cog.BasePredictor' assert hasattr(cls, 'setup'), f'setup method not found: {fullname}' diff --git a/python/cog/internal/openapi.json b/python/coglet/openapi.json similarity index 100% rename from python/cog/internal/openapi.json rename to python/coglet/openapi.json diff --git a/python/cog/internal/runner.py b/python/coglet/runner.py similarity index 97% rename from python/cog/internal/runner.py rename to python/coglet/runner.py index 5c23d15..c1dffef 100644 --- a/python/cog/internal/runner.py +++ b/python/coglet/runner.py @@ -5,8 +5,7 @@ import re from typing import Any, AsyncGenerator, Dict -import cog -from cog.internal import adt, util +from coglet import adt, api, util def _kwargs(adt_ins: Dict[str, adt.Input], inputs: Dict[str, Any]) -> Dict[str, Any]: @@ -96,8 +95,8 @@ async def setup(self) -> None: kwargs['weights'] = url self.predictor.setup(weights=url) elif os.path.exists(path): - kwargs['weights'] = cog.Path(path) - self.predictor.setup(weights=cog.Path(path)) + kwargs['weights'] = api.Path(path) + self.predictor.setup(weights=api.Path(path)) else: kwargs['weights'] = None if inspect.iscoroutinefunction(self.predictor.setup): diff --git a/python/cog/internal/schemas.py b/python/coglet/schemas.py similarity index 99% rename from python/cog/internal/schemas.py rename to python/coglet/schemas.py index 5c1fa5a..11dd260 100644 --- a/python/cog/internal/schemas.py +++ b/python/coglet/schemas.py @@ -2,7 +2,7 @@ import os.path from typing import Any, Dict, Optional, Union -from cog.internal import adt, util +from coglet import adt, util def _from_json_type(prop: Dict[str, Any]) -> adt.Type: diff --git a/python/cog/internal/util.py b/python/coglet/util.py similarity index 90% rename from python/cog/internal/util.py rename to python/coglet/util.py index 8e85d7f..a329394 100644 --- a/python/cog/internal/util.py +++ b/python/coglet/util.py @@ -2,8 +2,7 @@ from datetime import datetime, timezone from typing import Any, Tuple -import cog -from cog.internal import adt +from coglet import adt, api def check_cog_type(tpe: type) -> Tuple[adt.Type, bool]: @@ -47,9 +46,9 @@ def normalize_value(expected: adt.Type, value: Any) -> Any: if expected is adt.Type.FLOAT: return float(value) elif expected is adt.Type.PATH: - return cog.Path(value) if type(value) is str else value + return api.Path(value) if type(value) is str else value elif expected is adt.Type.SECRET: - return cog.Secret(value) if type(value) is str else value + return api.Secret(value) if type(value) is str else value else: return value diff --git a/python/pyproject.toml b/python/pyproject.toml index 17e9643..60b1f71 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -1,7 +1,7 @@ [project] -name = "coggo" +name = "coglet" version = "0.1.0" -description = "Add your description here" +description = "Minimum viable Cog runtime" readme = "README.md" requires-python = ">=3.9" classifiers = [ @@ -31,3 +31,6 @@ build-backend = "setuptools.build_meta" [tool.pytest.ini_options] asyncio_default_fixture_loop_scope = "function" +filterwarnings = [ + "ignore::ImportWarning", +] diff --git a/python/tests/test_file_runner.py b/python/tests/test_file_runner.py index 2d1c1a6..3ed5839 100644 --- a/python/tests/test_file_runner.py +++ b/python/tests/test_file_runner.py @@ -7,7 +7,7 @@ import time from typing import Dict, List, Optional -from cog.internal.file_runner import FileRunner +from coglet.file_runner import FileRunner def setup_signals() -> List[int]: @@ -28,7 +28,7 @@ def file_runner( cmd = [ sys.executable, '-m', - 'cog.internal.file_runner', + 'coglet', '--working-dir', tmp_path, '--module-name', diff --git a/python/tests/test_file_runner_async.py b/python/tests/test_file_runner_async.py index 8b49475..d338ff2 100644 --- a/python/tests/test_file_runner_async.py +++ b/python/tests/test_file_runner_async.py @@ -5,7 +5,7 @@ import pytest -from cog.internal.file_runner import FileRunner +from coglet.file_runner import FileRunner from .test_file_runner import file_runner, setup_signals diff --git a/python/tests/test_file_runner_iterator.py b/python/tests/test_file_runner_iterator.py index cde8916..46a5d0c 100644 --- a/python/tests/test_file_runner_iterator.py +++ b/python/tests/test_file_runner_iterator.py @@ -4,7 +4,7 @@ import time from typing import List, Optional -from cog.internal.file_runner import FileRunner +from coglet.file_runner import FileRunner from tests.test_file_runner import file_runner, setup_signals diff --git a/python/tests/test_predictors.py b/python/tests/test_predictors.py index 4109eec..2d323ae 100644 --- a/python/tests/test_predictors.py +++ b/python/tests/test_predictors.py @@ -6,8 +6,8 @@ import pytest -import cog -from cog.internal import inspector, runner, schemas +import coglet.api +from coglet import inspector, runner, schemas # Test predictors in tests/schemas # * run prediction with input/output fixture @@ -62,13 +62,13 @@ def test_schema(predictor): assert schemas.to_json_output(p) == schema['components']['schemas']['Output'] assert schemas.to_json_schema(p) == schema - eq = cog.Secret.__eq__ + eq = coglet.api.Secret.__eq__ if predictor == 'secrets': - cog.Secret.__eq__ = lambda self, other: type(other) is cog.Secret + coglet.api.Secret.__eq__ = lambda self, other: type(other) is coglet.api.Secret assert schemas.from_json_input(schema) == p.inputs assert schemas.from_json_output(schema) == p.output assert schemas.from_json_schema(module_name, class_name, schema) == p if predictor == 'secrets': - cog.Secret.__eq__ = eq + coglet.api.Secret.__eq__ = eq diff --git a/python/tests/test_util.py b/python/tests/test_util.py index 53db8d6..a07db10 100644 --- a/python/tests/test_util.py +++ b/python/tests/test_util.py @@ -1,7 +1,7 @@ from typing import List -from cog import Path, Secret -from cog.internal import adt, util +from coglet import adt, util +from coglet.api import Path, Secret def test_check_cog_type(): diff --git a/python/tests/test_weights.py b/python/tests/test_weights.py index 482a644..3a8beb6 100644 --- a/python/tests/test_weights.py +++ b/python/tests/test_weights.py @@ -2,7 +2,7 @@ import pytest -from cog.internal import inspector, runner +from coglet import inspector, runner @pytest.mark.asyncio