Merge branch 'main' into add-fast-generator

Signed-off-by: Will Sackfield <[email protected]>

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

python/cog/server/runner.py

@@ -355,22 +355,22 @@ 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 = []
 
         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:

python/tests/server/test_worker.py

@@ -134,21 +134,21 @@
     ),
     (
         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",
     ),
 ]
 
 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()

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) == "<class 'cog.predictor.Input'>"
-    ), "Predict input type should be the predictor Input."
+    assert str(input_type) == "<class 'cog.predictor.Input'>", (
+        "Predict input type should be the predictor Input."
+    )
     assert (
         str(output_type) == "<class 'cog.predictor.get_output_type.<locals>.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) == "<class 'cog.predictor.Input'>"
-        ), "Predict input type should be the predictor Input."
+        assert str(input_type) == "<class 'cog.predictor.Input'>", (
+            "Predict input type should be the predictor Input."
+        )
         assert (
             str(output_type)
             == "<class 'cog.predictor.get_output_type.<locals>.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) == "<class 'cog.predictor.Input'>"
-        ), "Predict input type should be the predictor Input."
+        assert str(input_type) == "<class 'cog.predictor.Input'>", (
+            "Predict input type should be the predictor Input."
+        )
         assert (
             str(output_type)
             == "<class 'cog.predictor.get_output_type.<locals>.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) == "<class 'cog.predictor.TrainingInput'>"
-        ), "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) == "<class 'cog.predictor.TrainingInput'>", (
+            "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) == "<class 'cog.predictor.TrainingInput'>"
-        ), "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) == "<class 'cog.predictor.TrainingInput'>", (
+            "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"

python/tests/test_wait.py

@@ -78,19 +78,19 @@ 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():
     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."
+        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."
+        )

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"

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

test-integration/test_integration/test_build.py

@@ -389,6 +389,23 @@ def test_pip_freeze(docker_image):
     )
 
 
+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
+
+
 def test_fast_build(docker_image):
     project_dir = Path(__file__).parent / "fixtures/fast-build"
     weights_file = os.path.join(project_dir, "weights.h5")