diff --git a/Makefile b/Makefile index 5f287822b3..068dc97e56 100644 --- a/Makefile +++ b/Makefile @@ -62,7 +62,7 @@ install: $(COG_BINARIES) .PHONY: clean clean: - rm -rf build dist pkg/dockerfile/embed + rm -rf .tox build dist pkg/dockerfile/embed rm -f $(COG_BINARIES) .PHONY: test-go diff --git a/python/cog/server/runner.py b/python/cog/server/runner.py index 1b752ae27a..a911a705ef 100644 --- a/python/cog/server/runner.py +++ b/python/cog/server/runner.py @@ -355,12 +355,12 @@ def wait(self, timeout: Optional[float] = None) -> None: self._fut.result(timeout=timeout) def set_output_type(self, *, multi: bool) -> None: - assert ( - self._output_type_multi is None - ), "Predictor unexpectedly returned multiple output types" - assert ( - self._p.output is None - ), "Predictor unexpectedly returned output type after output" + assert self._output_type_multi is None, ( + "Predictor unexpectedly returned multiple output types" + ) + assert self._p.output is None, ( + "Predictor unexpectedly returned output type after output" + ) if multi: self._p.output = [] @@ -368,9 +368,9 @@ def set_output_type(self, *, multi: bool) -> None: self._output_type_multi = multi def append_output(self, output: Any) -> None: - assert ( - self._output_type_multi is not None - ), "Predictor unexpectedly returned output before output type" + assert self._output_type_multi is not None, ( + "Predictor unexpectedly returned output before output type" + ) uploaded_output = self._upload_files(output) if self._output_type_multi: diff --git a/python/tests/server/test_worker.py b/python/tests/server/test_worker.py index da40d81287..dd1b4e82fb 100644 --- a/python/tests/server/test_worker.py +++ b/python/tests/server/test_worker.py @@ -134,7 +134,7 @@ ), ( WorkerConfig("logging_async", setup=False, min_python=(3, 11), is_async=True), - ("writing to stdout at import time\n" "setting up predictor\n"), + "writing to stdout at import time\nsetting up predictor\n", "writing to stderr at import time\n", ), ] @@ -142,13 +142,13 @@ PREDICT_LOGS_FIXTURES = [ ( WorkerConfig("logging"), - ("writing from C\n" "writing with print\n"), - ("WARNING:root:writing log message\n" "writing to stderr\n"), + "writing from C\nwriting with print\n", + "WARNING:root:writing log message\nwriting to stderr\n", ), ( WorkerConfig("logging_async", min_python=(3, 11), is_async=True), - ("writing with print\n"), - ("WARNING:root:writing log message\n" "writing to stderr\n"), + "writing with print\n", + "WARNING:root:writing log message\nwriting to stderr\n", ), ] @@ -214,14 +214,14 @@ def handle_event(self, event: _PublicEventType): if self.output_type.multi: self.output.append(event.payload) else: - assert ( - self.output is None - ), "Should not get multiple outputs for output type single" + assert self.output is None, ( + "Should not get multiple outputs for output type single" + ) self.output = event.payload elif isinstance(event, PredictionOutputType): - assert ( - self.output_type is None - ), "Should not get multiple output type events" + assert self.output_type is None, ( + "Should not get multiple output type events" + ) self.output_type = event if self.output_type.multi: self.output = [] @@ -863,7 +863,7 @@ def simulate_predict_success(self, state: PredictState): events.append(PredictionOutputType(multi=True)) for i in range(steps): events.append( - PredictionOutput(payload=f"NAME={name},STEP={i+1}"), + PredictionOutput(payload=f"NAME={name},STEP={i + 1}"), ) events.append(Done(canceled=state.canceled)) @@ -909,7 +909,7 @@ def await_predict(self, state: PredictState): assert state.result.output == f"NAME={name}" else: assert state.result.output == [ - f"NAME={name},STEP={i+1}" for i in range(steps) + f"NAME={name},STEP={i + 1}" for i in range(steps) ] assert state.result.done == Done() diff --git a/python/tests/test_config.py b/python/tests/test_config.py index d594920d6a..a60db1dca1 100644 --- a/python/tests/test_config.py +++ b/python/tests/test_config.py @@ -21,9 +21,9 @@ def test_predictor_predict_ref_env_var(): config = Config() config_predict_ref = config.predictor_predict_ref del os.environ[COG_PREDICT_TYPE_STUB_ENV_VAR] - assert ( - config_predict_ref == predict_ref - ), "Predict Reference should come from the environment variable." + assert config_predict_ref == predict_ref, ( + "Predict Reference should come from the environment variable." + ) def test_predictor_predict_ref_no_env_var(): @@ -40,9 +40,9 @@ def test_predictor_predict_ref_no_env_var(): """) config = Config() config_predict_ref = config.predictor_predict_ref - assert ( - config_predict_ref == "predict.py:Predictor" - ), "Predict Reference should come from the cog config file." + assert config_predict_ref == "predict.py:Predictor", ( + "Predict Reference should come from the cog config file." + ) os.chdir(pwd) @@ -59,9 +59,9 @@ def test_config_initial_values(): del os.environ[COG_PREDICT_TYPE_STUB_ENV_VAR] config = Config(config={"predict": "predict.py:Predictor"}) config_predict_ref = config.predictor_predict_ref - assert ( - config_predict_ref == "predict.py:Predictor" - ), "Predict Reference should come from the initial config dictionary." + assert config_predict_ref == "predict.py:Predictor", ( + "Predict Reference should come from the initial config dictionary." + ) def test_predictor_train_ref_env_var(): @@ -70,9 +70,9 @@ def test_predictor_train_ref_env_var(): config = Config() config_train_ref = config.predictor_train_ref del os.environ[COG_TRAIN_TYPE_STUB_ENV_VAR] - assert ( - config_train_ref == train_ref - ), "Train Reference should come from the environment variable." + assert config_train_ref == train_ref, ( + "Train Reference should come from the environment variable." + ) def test_predictor_train_ref_no_env_var(): @@ -81,9 +81,9 @@ def test_predictor_train_ref_no_env_var(): del os.environ[COG_TRAIN_TYPE_STUB_ENV_VAR] config = Config(config={"train": train_ref}) config_train_ref = config.predictor_train_ref - assert ( - config_train_ref == train_ref - ), "Train Reference should come from the initial config dictionary." + assert config_train_ref == train_ref, ( + "Train Reference should come from the initial config dictionary." + ) def test_requires_gpu_env_var(): @@ -100,27 +100,27 @@ def test_requires_gpu_no_env_var(): del os.environ[COG_GPU_ENV_VAR] config = Config(config={"build": {"gpu": False}}) config_gpu = config.requires_gpu - assert ( - not config_gpu - ), "Requires GPU should come from the initial config dictionary." + assert not config_gpu, ( + "Requires GPU should come from the initial config dictionary." + ) def test_get_predictor_ref_predict(): train_ref = "predict.py:Predictor" config = Config(config={"train": train_ref}) config_train_ref = config.get_predictor_ref(Mode.TRAIN) - assert ( - train_ref == config_train_ref - ), "The train ref should equal the config train ref." + assert train_ref == config_train_ref, ( + "The train ref should equal the config train ref." + ) def test_get_predictor_ref_train(): predict_ref = "predict.py:Predictor" config = Config(config={"predict": predict_ref}) config_predict_ref = config.get_predictor_ref(Mode.PREDICT) - assert ( - predict_ref == config_predict_ref - ), "The predict ref should equal the config predict ref." + assert predict_ref == config_predict_ref, ( + "The predict ref should equal the config predict ref." + ) def test_get_predictor_types_with_env_var(): @@ -144,9 +144,9 @@ def predict(self, msg: str) -> ModelOutput: input_type, output_type, is_async = config.get_predictor_types(Mode.PREDICT) del os.environ[COG_PREDICT_CODE_STRIP_ENV_VAR] del os.environ[COG_PREDICT_TYPE_STUB_ENV_VAR] - assert ( - str(input_type) == "" - ), "Predict input type should be the predictor Input." + assert str(input_type) == "", ( + "Predict input type should be the predictor Input." + ) assert ( str(output_type) == ".Output'>" ), "Predict output type should be the predictor Output." @@ -179,9 +179,9 @@ def predict(self, msg: str) -> ModelOutput: predict_ref = f"{predict_python_file}:Predictor" config = Config(config={"predict": predict_ref}) input_type, output_type, is_async = config.get_predictor_types(Mode.PREDICT) - assert ( - str(input_type) == "" - ), "Predict input type should be the predictor Input." + assert str(input_type) == "", ( + "Predict input type should be the predictor Input." + ) assert ( str(output_type) == ".Output'>" @@ -215,9 +215,9 @@ async def predict(self, msg: str) -> ModelOutput: predict_ref = f"{predict_python_file}:Predictor" config = Config(config={"predict": predict_ref}) input_type, output_type, is_async = config.get_predictor_types(Mode.PREDICT) - assert ( - str(input_type) == "" - ), "Predict input type should be the predictor Input." + assert str(input_type) == "", ( + "Predict input type should be the predictor Input." + ) assert ( str(output_type) == ".Output'>" @@ -249,12 +249,12 @@ def train( train_ref = f"{predict_python_file}:train" config = Config(config={"train": train_ref}) input_type, output_type, is_async = config.get_predictor_types(Mode.TRAIN) - assert ( - str(input_type) == "" - ), "Predict input type should be the training Input." - assert str(output_type).endswith( - "TrainingOutput'>" - ), "Predict output type should be the training Output." + assert str(input_type) == "", ( + "Predict input type should be the training Input." + ) + assert str(output_type).endswith("TrainingOutput'>"), ( + "Predict output type should be the training Output." + ) assert not is_async, "is_async should be False for normal functions" @@ -282,10 +282,10 @@ async def train( train_ref = f"{predict_python_file}:train" config = Config(config={"train": train_ref}) input_type, output_type, is_async = config.get_predictor_types(Mode.TRAIN) - assert ( - str(input_type) == "" - ), "Predict input type should be the training Input." - assert str(output_type).endswith( - "TrainingOutput'>" - ), "Predict output type should be the training Output." + assert str(input_type) == "", ( + "Predict input type should be the training Input." + ) + assert str(output_type).endswith("TrainingOutput'>"), ( + "Predict output type should be the training Output." + ) assert is_async, "is_async should be True for async functions" diff --git a/python/tests/test_wait.py b/python/tests/test_wait.py index 07afb9266d..aedf473a03 100644 --- a/python/tests/test_wait.py +++ b/python/tests/test_wait.py @@ -78,9 +78,9 @@ def test_wait_for_env_no_env_vars(): 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." + assert result, ( + "We should return true if we have no env vars associated with the wait." + ) def test_wait_for_env(): @@ -88,9 +88,9 @@ def test_wait_for_env(): 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." + 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] @@ -118,6 +118,6 @@ def test_wait_inserts_pythonpath(): 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." + assert expected_path == current_python_path, ( + "Our python path should be updated with the pyenv path." + ) diff --git a/test-integration/test_integration/fixtures/apt-packages/cog.yaml b/test-integration/test_integration/fixtures/apt-packages/cog.yaml new file mode 100644 index 0000000000..51fab884ff --- /dev/null +++ b/test-integration/test_integration/fixtures/apt-packages/cog.yaml @@ -0,0 +1,8 @@ +build: + gpu: true + python_version: "3.9" + system_packages: + - "git" + run: + - command: git --version +predict: "predict.py:Predictor" diff --git a/test-integration/test_integration/fixtures/apt-packages/predict.py b/test-integration/test_integration/fixtures/apt-packages/predict.py new file mode 100644 index 0000000000..44f6992b01 --- /dev/null +++ b/test-integration/test_integration/fixtures/apt-packages/predict.py @@ -0,0 +1,6 @@ +from cog import BasePredictor + + +class Predictor(BasePredictor): + def predict(self, s: str) -> str: + return "hello " + s diff --git a/test-integration/test_integration/test_build.py b/test-integration/test_integration/test_build.py index 7c76d0c19b..e5a9c93df3 100644 --- a/test-integration/test_integration/test_build.py +++ b/test-integration/test_integration/test_build.py @@ -387,3 +387,20 @@ def test_pip_freeze(docker_image): pip_freeze == "anyio==4.4.0\nattrs==23.2.0\ncertifi==2024.8.30\ncharset-normalizer==3.3.2\nclick==8.1.7\nexceptiongroup==1.2.2\nfastapi==0.98.0\nh11==0.14.0\nhttptools==0.6.1\nidna==3.8\npydantic==1.10.18\npython-dotenv==1.0.1\nPyYAML==6.0.2\nrequests==2.32.3\nsniffio==1.3.1\nstarlette==0.27.0\nstructlog==24.4.0\ntyping_extensions==4.12.2\nurllib3==2.2.2\nuvicorn==0.30.6\nuvloop==0.20.0\nwatchfiles==0.24.0\nwebsockets==13.0.1\n" ) + + +def test_cog_installs_apt_packages(docker_image): + project_dir = Path(__file__).parent / "fixtures/apt-packages" + build_process = subprocess.run( + [ + "cog", + "build", + "-t", + docker_image, + ], + cwd=project_dir, + capture_output=True, + ) + # Test that the build completes successfully. + # If the apt-packages weren't installed the run command would fail. + assert build_process.returncode == 0