Repository: replicate/cog

Issue:
`cog predict` should print JSON and plain text output to stdout by default
If you don't pass `-o`, it should just print the output to stdout. This is a regression. It was working at some point.


Discussion:

Resolving patch:
diff --git a/pkg/docker/cog.py b/pkg/docker/cog.py
--- a/pkg/docker/cog.py
+++ b/pkg/docker/cog.py
@@ -39,7 +39,7 @@ def setup(self):
         pass
 
     @abstractmethod
-    def run(self, **kwargs):
+    def predict(self, **kwargs):
         pass
 
 
@@ -53,7 +53,8 @@ def make_app(self) -> Flask:
         app = Flask(__name__)
         setup_time = time.time() - start_time
 
-        @app.route("/infer", methods=["POST"])
+        @app.route("/predict", methods=["POST"])
+        @app.route("/infer", methods=["POST"])  # deprecated
         def handle_request():
             start_time = time.time()
 
@@ -69,7 +70,7 @@ def handle_request():
                         )
                     raw_inputs[key] = val
 
-                if hasattr(self.model.run, "_inputs"):
+                if hasattr(self.model.predict, "_inputs"):
                     try:
                         inputs = validate_and_convert_inputs(
                             self.model, raw_inputs, cleanup_functions
@@ -96,8 +97,8 @@ def ping():
         @app.route("/help")
         def help():
             args = {}
-            if hasattr(self.model.run, "_inputs"):
-                input_specs = self.model.run._inputs
+            if hasattr(self.model.predict, "_inputs"):
+                input_specs = self.model.predict._inputs
                 for name, spec in input_specs.items():
                     arg: Dict[str, Any] = {
                         "type": _type_name(spec.type),
@@ -185,8 +186,8 @@ def ping():
         @app.route("/help")
         def help():
             args = {}
-            if hasattr(self.model.run, "_inputs"):
-                input_specs = self.model.run._inputs
+            if hasattr(self.model.predict, "_inputs"):
+                input_specs = self.model.predict._inputs
                 for name, spec in input_specs.items():
                     arg = {
                         "type": _type_name(spec.type),
@@ -245,7 +246,7 @@ def __init__(
         )
         self.should_exit = False
         self.setup_time_queue = input_queue + "-setup-time"
-        self.infer_time_queue = input_queue + "-run-time"
+        self.predict_time_queue = input_queue + "-run-time"
         self.stats_queue_length = 100
 
         sys.stderr.write(
@@ -311,9 +312,11 @@ def start(self):
                     continue
 
                 message = json.loads(message_json)
-                infer_id = message["id"]
+                prediction_id = message["id"]
                 response_queue = message["response_queue"]
-                sys.stderr.write(f"Received message {infer_id} on {self.input_queue}\n")
+                sys.stderr.write(
+                    f"Received message {prediction_id} on {self.input_queue}\n"
+                )
                 cleanup_functions = []
                 try:
                     start_time = time.time()
@@ -324,7 +327,7 @@ def start(self):
                     )  # xdel to be able to get stream size
                     run_time = time.time() - start_time
                     self.redis.xadd(
-                        self.infer_time_queue,
+                        self.predict_time_queue,
                         fields={"duration": run_time},
                         maxlen=self.stats_queue_length,
                     )
@@ -417,7 +420,7 @@ def upload_to_temp(self, path: Path) -> str:
 def validate_and_convert_inputs(
     model: Model, raw_inputs: Dict[str, Any], cleanup_functions: List[Callable]
 ) -> Dict[str, Any]:
-    input_specs = model.run._inputs
+    input_specs = model.predict._inputs
     inputs = {}
 
     for name, input_spec in input_specs.items():
@@ -513,7 +516,7 @@ def run_model(model, inputs, cleanup_functions):
     Run the model on the inputs, and append resulting paths
     to cleanup functions for removal.
     """
-    result = model.run(**inputs)
+    result = model.predict(**inputs)
     if isinstance(result, Path):
         cleanup_functions.append(result.unlink)
     return result


Test patch:
diff --git a/end-to-end-test/end_to_end_test/conftest.py b/end-to-end-test/end_to_end_test/conftest.py
--- a/end-to-end-test/end_to_end_test/conftest.py
+++ b/end-to-end-test/end_to_end_test/conftest.py
@@ -55,7 +55,7 @@ def cog_server(registry_host) -> Generator[CogServer, None, None]:
 @pytest.fixture
 def project_dir(tmpdir_factory):
     tmpdir = tmpdir_factory.mktemp("project")
-    with open(tmpdir / "infer.py", "w") as f:
+    with open(tmpdir / "predict.py", "w") as f:
         f.write(
             """
 import time
@@ -70,7 +70,7 @@ def setup(self):
     @cog.input("text", type=str)
     @cog.input("path", type=Path)
     @cog.input("output_file", type=bool, default=False)
-    def run(self, text, path, output_file):
+    def predict(self, text, path, output_file):
         time.sleep(1)
         with open(path) as f:
             output = self.foo + text + f.read()
@@ -87,7 +87,7 @@ def run(self, text, path, output_file):
     with open(tmpdir / "cog.yaml", "w") as f:
         cog_yaml = """
 name: andreas/hello-world
-model: infer.py:Model
+model: predict.py:Model
 examples:
   - input:
       text: "foo"
diff --git a/end-to-end-test/end_to_end_test/test_queue_worker.py b/end-to-end-test/end_to_end_test/test_queue_worker.py
--- a/end-to-end-test/end_to_end_test/test_queue_worker.py
+++ b/end-to-end-test/end_to_end_test/test_queue_worker.py
@@ -32,7 +32,7 @@ def test_queue_worker(cog_server, project_dir, redis_port, tmpdir_factory):
     upload_url = f"http://{local_ip}:{controller_port}/upload"
     redis_host = local_ip
     worker_name = "test-worker"
-    infer_queue_name = "infer-queue"
+    predict_queue_name = "predict-queue"
     response_queue_name = "response-queue"
 
     wait_for_port(redis_host, redis_port)
@@ -46,22 +46,22 @@ def test_queue_worker(cog_server, project_dir, redis_port, tmpdir_factory):
             "cog-redis-queue-worker",
             redis_host,
             str(redis_port),
-            infer_queue_name,
+            predict_queue_name,
             upload_url,
             worker_name,
         ],
     ):
         redis_client.xgroup_create(
-            mkstream=True, groupname=infer_queue_name, name=infer_queue_name, id="$"
+            mkstream=True, groupname=predict_queue_name, name=predict_queue_name, id="$"
         )
 
-        infer_id = random_string(10)
+        predict_id = random_string(10)
         redis_client.xadd(
-            name=infer_queue_name,
+            name=predict_queue_name,
             fields={
                 "value": json.dumps(
                     {
-                        "id": infer_id,
+                        "id": predict_id,
                         "inputs": {
                             "text": {"value": "bar"},
                             "path": {
@@ -80,13 +80,13 @@ def test_queue_worker(cog_server, project_dir, redis_port, tmpdir_factory):
         response = json.loads(redis_client.brpop(response_queue_name)[1])["value"]
         assert response == "foobartest"
 
-        infer_id = random_string(10)
+        predict_id = random_string(10)
         redis_client.xadd(
-            name=infer_queue_name,
+            name=predict_queue_name,
             fields={
                 "value": json.dumps(
                     {
-                        "id": infer_id,
+                        "id": predict_id,
                         "inputs": {
                             "text": {"value": "bar"},
                             "output_file": {"value": "true"},
diff --git a/end-to-end-test/end_to_end_test/test_server.py b/end-to-end-test/end_to_end_test/test_server.py
--- a/end-to-end-test/end_to_end_test/test_server.py
+++ b/end-to-end-test/end_to_end_test/test_server.py
@@ -101,7 +101,7 @@ def test_server_end_to_end(cog_server, project_dir, tmpdir_factory):
     ).communicate()
     paths = sorted(glob(str(download_dir / "*.*")))
     filenames = [os.path.basename(f) for f in paths]
-    assert filenames == ["cog.yaml", "infer.py", "myfile.txt"]
+    assert filenames == ["cog.yaml", "myfile.txt", "predict.py"]
 
     with open(download_dir / "cog-example-output/output.02.txt") as f:
         assert f.read() == "fooquxbaz"
@@ -124,7 +124,7 @@ def test_server_end_to_end(cog_server, project_dir, tmpdir_factory):
             "cog",
             "--model",
             model_url,
-            "infer",
+            "predict",
             "-o",
             out_path,
             "-i",
diff --git a/pkg/docker/cog_test.py b/pkg/docker/cog_test.py
--- a/pkg/docker/cog_test.py
+++ b/pkg/docker/cog_test.py
@@ -26,11 +26,11 @@ class Model(cog.Model):
         def setup(self):
             self.foo = "foo"
 
-        def run(self):
+        def predict(self):
             return self.foo + "bar"
 
     client = make_client(Model())
-    resp = client.post("/infer")
+    resp = client.post("/predict")
     assert resp.status_code == 200
     assert resp.data == b"foobar"
 
@@ -41,11 +41,11 @@ def setup(self):
             self.foo = "foo"
 
         @cog.input("text", type=str)
-        def run(self, text):
+        def predict(self, text):
             return self.foo + text
 
     client = make_client(Model())
-    resp = client.post("/infer", data={"text": "baz"})
+    resp = client.post("/predict", data={"text": "baz"})
     assert resp.status_code == 200
     assert resp.data == b"foobaz"
 
@@ -56,11 +56,11 @@ def setup(self):
             self.foo = "foo"
 
         @cog.input("text", type=str)
-        def run(self, text):
+        def predict(self, text):
             return self.foo + text
 
     client = make_client(Model())
-    resp = client.post("/infer", data={"text": "baz", "text2": "qux"})
+    resp = client.post("/predict", data={"text": "baz", "text2": "qux"})
     assert resp.status_code == 400
 
 
@@ -73,7 +73,7 @@ def setup(self):
                 self.foo = "foo"
 
             @cog.input("text", type=str)
-            def run(self, bad):
+            def predict(self, bad):
                 return self.foo + "bar"
 
 
@@ -83,15 +83,15 @@ def setup(self):
             self.foo = "foo"
 
         @cog.input("num", type=int)
-        def run(self, num):
+        def predict(self, num):
             num2 = num ** 3
             return self.foo + str(num2)
 
     client = make_client(Model())
-    resp = client.post("/infer", data={"num": 3})
+    resp = client.post("/predict", data={"num": 3})
     assert resp.status_code == 200
     assert resp.data == b"foo27"
-    resp = client.post("/infer", data={"num": -3})
+    resp = client.post("/predict", data={"num": -3})
     assert resp.status_code == 200
     assert resp.data == b"foo-27"
 
@@ -102,12 +102,12 @@ def setup(self):
             self.foo = "foo"
 
         @cog.input("num", type=int)
-        def run(self, num):
+        def predict(self, num):
             num2 = num ** 2
             return self.foo + str(num2)
 
     client = make_client(Model())
-    resp = client.post("/infer", data={"num": "foo"})
+    resp = client.post("/predict", data={"num": "foo"})
     assert resp.status_code == 400
 
 
@@ -117,15 +117,15 @@ def setup(self):
             self.foo = "foo"
 
         @cog.input("num", type=int, default=5)
-        def run(self, num):
+        def predict(self, num):
             num2 = num ** 2
             return self.foo + str(num2)
 
     client = make_client(Model())
-    resp = client.post("/infer", data={"num": 3})
+    resp = client.post("/predict", data={"num": 3})
     assert resp.status_code == 200
     assert resp.data == b"foo9"
-    resp = client.post("/infer")
+    resp = client.post("/predict")
     assert resp.status_code == 200
     assert resp.data == b"foo25"
 
@@ -136,18 +136,18 @@ def setup(self):
             self.foo = "foo"
 
         @cog.input("num", type=float)
-        def run(self, num):
+        def predict(self, num):
             num2 = num ** 3
             return self.foo + str(num2)
 
     client = make_client(Model())
-    resp = client.post("/infer", data={"num": 3})
+    resp = client.post("/predict", data={"num": 3})
     assert resp.status_code == 200
     assert resp.data == b"foo27.0"
-    resp = client.post("/infer", data={"num": 3.5})
+    resp = client.post("/predict", data={"num": 3.5})
     assert resp.status_code == 200
     assert resp.data == b"foo42.875"
-    resp = client.post("/infer", data={"num": -3.5})
+    resp = client.post("/predict", data={"num": -3.5})
     assert resp.status_code == 200
     assert resp.data == b"foo-42.875"
 
@@ -158,12 +158,12 @@ def setup(self):
             self.foo = "foo"
 
         @cog.input("num", type=float)
-        def run(self, num):
+        def predict(self, num):
             num2 = num ** 2
             return self.foo + str(num2)
 
     client = make_client(Model())
-    resp = client.post("/infer", data={"num": "foo"})
+    resp = client.post("/predict", data={"num": "foo"})
     assert resp.status_code == 400
 
 
@@ -173,17 +173,17 @@ def setup(self):
             self.foo = "foo"
 
         @cog.input("flag", type=bool)
-        def run(self, flag):
+        def predict(self, flag):
             if flag:
                 return self.foo + "yes"
             else:
                 return self.foo + "no"
 
     client = make_client(Model())
-    resp = client.post("/infer", data={"flag": True})
+    resp = client.post("/predict", data={"flag": True})
     assert resp.status_code == 200
     assert resp.data == b"fooyes"
-    resp = client.post("/infer", data={"flag": False})
+    resp = client.post("/predict", data={"flag": False})
     assert resp.status_code == 200
     assert resp.data == b"foono"
 
@@ -194,14 +194,14 @@ def setup(self):
             self.foo = "foo"
 
         @cog.input("flag", type=bool)
-        def run(self, flag):
+        def predict(self, flag):
             if flag:
                 return self.foo + "yes"
             else:
                 return self.foo + "no"
 
     client = make_client(Model())
-    resp = client.post("/infer", data={"flag": "foo"})
+    resp = client.post("/predict", data={"flag": "foo"})
     assert resp.status_code == 400
 
 
@@ -213,18 +213,18 @@ def setup(self):
         @cog.input("num1", type=float, min=3, max=10.5)
         @cog.input("num2", type=float, min=-4)
         @cog.input("num3", type=int, max=-4)
-        def run(self, num1, num2, num3):
+        def predict(self, num1, num2, num3):
             return num1 + num2 + num3
 
     client = make_client(Model())
-    resp = client.post("/infer", data={"num1": 3, "num2": -4, "num3": -4})
+    resp = client.post("/predict", data={"num1": 3, "num2": -4, "num3": -4})
     assert resp.status_code == 200
     assert resp.data == b"-5.0\n"
-    resp = client.post("/infer", data={"num1": 2, "num2": -4, "num3": -4})
+    resp = client.post("/predict", data={"num1": 2, "num2": -4, "num3": -4})
     assert resp.status_code == 400
-    resp = client.post("/infer", data={"num1": 3, "num2": -4.1, "num3": -4})
+    resp = client.post("/predict", data={"num1": 3, "num2": -4.1, "num3": -4})
     assert resp.status_code == 400
-    resp = client.post("/infer", data={"num1": 3, "num2": -4, "num3": -3})
+    resp = client.post("/predict", data={"num1": 3, "num2": -4, "num3": -3})
     assert resp.status_code == 400
 
 
@@ -235,11 +235,11 @@ def setup(self):
 
         @cog.input("text", type=str, options=["foo", "bar"])
         @cog.input("num", type=int, options=[1, 2, 3])
-        def run(self, text, num):
+        def predict(self, text, num):
             return text + ("a" * num)
 
     client = make_client(Model())
-    resp = client.post("/infer", data={"text": "foo", "num": 2})
+    resp = client.post("/predict", data={"text": "foo", "num": 2})
     assert resp.status_code == 200
     assert resp.data == b"fooaa"
 
@@ -252,7 +252,7 @@ def setup(self):
                 pass
 
             @cog.input("text", type=str, options=[])
-            def run(self, text):
+            def predict(self, text):
                 return text
 
     with pytest.raises(ValueError):
@@ -262,7 +262,7 @@ def setup(self):
                 pass
 
             @cog.input("text", type=str, options=["foo"])
-            def run(self, text):
+            def predict(self, text):
                 return text
 
     with pytest.raises(ValueError):
@@ -272,7 +272,7 @@ def setup(self):
                 pass
 
             @cog.input("text", type=Path, options=["foo"])
-            def run(self, text):
+            def predict(self, text):
                 return text
 
 
@@ -283,13 +283,13 @@ def setup(self):
 
         @cog.input("text", type=str, options=["foo", "bar"])
         @cog.input("num", type=int, options=[1, 2, 3])
-        def run(self, text, num):
+        def predict(self, text, num):
             return text + ("a" * num)
 
     client = make_client(Model())
-    resp = client.post("/infer", data={"text": "baz", "num": 2})
+    resp = client.post("/predict", data={"text": "baz", "num": 2})
     assert resp.status_code == 400
-    resp = client.post("/infer", data={"text": "bar", "num": 4})
+    resp = client.post("/predict", data={"text": "bar", "num": 4})
     assert resp.status_code == 400
 
 
@@ -299,14 +299,14 @@ def setup(self):
             self.foo = "foo"
 
         @cog.input("path", type=Path)
-        def run(self, path):
+        def predict(self, path):
             with open(path) as f:
                 return self.foo + " " + f.read() + " " + os.path.basename(path)
 
     client = make_client(Model())
     path_data = (io.BytesIO(b"bar"), "foo.txt")
     resp = client.post(
-        "/infer", data={"path": path_data}, content_type="multipart/form-data"
+        "/predict", data={"path": path_data}, content_type="multipart/form-data"
     )
     assert resp.status_code == 200
     assert resp.data == b"foo bar foo.txt"
@@ -318,12 +318,12 @@ def setup(self):
             self.foo = "foo"
 
         @cog.input("path", type=Path)
-        def run(self, path):
+        def predict(self, path):
             with open(path) as f:
                 return self.foo + " " + f.read() + " " + os.path.basename(path)
 
     client = make_client(Model())
-    resp = client.post("/infer", data={"path": "bar"})
+    resp = client.post("/predict", data={"path": "bar"})
     assert resp.status_code == 400
 
 
@@ -333,7 +333,7 @@ def setup(self):
             self.foo = "foo"
 
         @cog.input("path", type=Path, default=None)
-        def run(self, path):
+        def predict(self, path):
             if path is None:
                 return "noneee"
             with open(path) as f:
@@ -342,11 +342,11 @@ def run(self, path):
     client = make_client(Model())
     path_data = (io.BytesIO(b"bar"), "foo.txt")
     resp = client.post(
-        "/infer", data={"path": path_data}, content_type="multipart/form-data"
+        "/predict", data={"path": path_data}, content_type="multipart/form-data"
     )
     assert resp.status_code == 200
     assert resp.data == b"foo bar foo.txt"
-    resp = client.post("/infer", data={})
+    resp = client.post("/predict", data={})
     assert resp.status_code == 200
     assert resp.data == b"noneee"
 
@@ -357,7 +357,7 @@ def setup(self):
             self.foo = "foo"
 
         @cog.input("text", type=str)
-        def run(self, text):
+        def predict(self, text):
             temp_dir = tempfile.mkdtemp()
             temp_path = os.path.join(temp_dir, "my_file.txt")
             with open(temp_path, "w") as f:
@@ -365,7 +365,7 @@ def run(self, text):
             return Path(temp_path)
 
     client = make_client(Model())
-    resp = client.post("/infer", data={"text": "baz"})
+    resp = client.post("/predict", data={"text": "baz"})
     assert resp.status_code == 200
     assert resp.content_type == "text/plain; charset=utf-8"
     assert resp.data == b"foobaz"
@@ -376,7 +376,7 @@ class Model(cog.Model):
         def setup(self):
             pass
 
-        def run(self):
+        def predict(self):
             temp_dir = tempfile.mkdtemp()
             temp_path = os.path.join(temp_dir, "my_file.bmp")
             img = Image.new("RGB", (255, 255), "red")
@@ -384,7 +384,7 @@ def run(self):
             return Path(temp_path)
 
     client = make_client(Model())
-    resp = client.post("/infer")
+    resp = client.post("/predict")
     assert resp.status_code == 200
     # need both image/bmp and image/x-ms-bmp until https://bugs.python.org/issue44211 is fixed
     assert resp.content_type in ["image/bmp", "image/x-ms-bmp"]
@@ -400,7 +400,7 @@ def setup(self):
         @cog.input("num1", type=int)
         @cog.input("num2", type=int, default=10)
         @cog.input("path", type=Path)
-        def run(self, text, num1, num2, path):
+        def predict(self, text, num1, num2, path):
             with open(path) as f:
                 path_contents = f.read()
             return self.foo + " " + text + " " + str(num1 * num2) + " " + path_contents
@@ -408,7 +408,7 @@ def run(self, text, num1, num2, path):
     client = make_client(Model())
     path_data = (io.BytesIO(b"bar"), "foo.txt")
     resp = client.post(
-        "/infer",
+        "/predict",
         data={"text": "baz", "num1": 5, "path": path_data},
         content_type="multipart/form-data",
     )
@@ -425,7 +425,7 @@ def setup(self):
         @cog.input("num1", type=int, help="First number")
         @cog.input("num2", type=int, default=10, help="Second number")
         @cog.input("path", type=Path, help="A file path")
-        def run(self, text, num1, num2, path):
+        def predict(self, text, num1, num2, path):
             with open(path) as f:
                 path_contents = f.read()
             return self.foo + " " + text + " " + str(num1 * num2) + " " + path_contents
@@ -461,7 +461,7 @@ class ModelSlow(cog.Model):
         def setup(self):
             time.sleep(0.5)
 
-        def run(self):
+        def predict(self):
             time.sleep(0.5)
             return ""
 
@@ -469,17 +469,17 @@ class ModelFast(cog.Model):
         def setup(self):
             pass
 
-        def run(self):
+        def predict(self):
             return ""
 
     client = make_client(ModelSlow())
-    resp = client.post("/infer")
+    resp = client.post("/predict")
     assert resp.status_code == 200
     assert 0.5 < float(resp.headers["X-Setup-Time"]) < 1.0
     assert 0.5 < float(resp.headers["X-Run-Time"]) < 1.0
 
     client = make_client(ModelFast())
-    resp = client.post("/infer")
+    resp = client.post("/predict")
     assert resp.status_code == 200
     assert float(resp.headers["X-Setup-Time"]) < 0.5
     assert float(resp.headers["X-Run-Time"]) < 0.5
diff --git a/pkg/docker/generate_test.go b/pkg/docker/generate_test.go
--- a/pkg/docker/generate_test.go
+++ b/pkg/docker/generate_test.go
@@ -56,7 +56,7 @@ func TestGenerateEmpty(t *testing.T) {
 	require.NoError(t, err)
 
 	conf, err := model.ConfigFromYAML([]byte(`
-model: infer.py:Model
+model: predict.py:Model
 `))
 	require.NoError(t, err)
 	require.NoError(t, conf.ValidateAndCompleteConfig())
@@ -106,7 +106,7 @@ environment:
   system_packages:
     - ffmpeg
     - cowsay
-model: infer.py:Model
+model: predict.py:Model
 `))
 	require.NoError(t, err)
 	require.NoError(t, conf.ValidateAndCompleteConfig())
@@ -159,10 +159,10 @@ CMD /usr/bin/cog-http-server`
 
 func testHelperScripts() string {
 	return `
-RUN echo '#!/usr/bin/env python\nimport sys\nimport cog\nimport os\nos.chdir("/code")\nsys.path.append("/code")\nfrom infer import Model\ncog.HTTPServer(Model()).start_server()' > /usr/bin/cog-http-server
+RUN echo '#!/usr/bin/env python\nimport sys\nimport cog\nimport os\nos.chdir("/code")\nsys.path.append("/code")\nfrom predict import Model\ncog.HTTPServer(Model()).start_server()' > /usr/bin/cog-http-server
 RUN chmod +x /usr/bin/cog-http-server
-RUN echo '#!/usr/bin/env python\nimport sys\nimport cog\nimport os\nos.chdir("/code")\nsys.path.append("/code")\nfrom infer import Model\ncog.AIPlatformPredictionServer(Model()).start_server()' > /usr/bin/cog-ai-platform-prediction-server
+RUN echo '#!/usr/bin/env python\nimport sys\nimport cog\nimport os\nos.chdir("/code")\nsys.path.append("/code")\nfrom predict import Model\ncog.AIPlatformPredictionServer(Model()).start_server()' > /usr/bin/cog-ai-platform-prediction-server
 RUN chmod +x /usr/bin/cog-ai-platform-prediction-server
-RUN echo '#!/usr/bin/env python\nimport sys\nimport cog\nimport os\nos.chdir("/code")\nsys.path.append("/code")\nfrom infer import Model\ncog.RedisQueueWorker(Model(), redis_host=sys.argv[1], redis_port=sys.argv[2], input_queue=sys.argv[3], upload_url=sys.argv[4], consumer_id=sys.argv[5]).start()' > /usr/bin/cog-redis-queue-worker
+RUN echo '#!/usr/bin/env python\nimport sys\nimport cog\nimport os\nos.chdir("/code")\nsys.path.append("/code")\nfrom predict import Model\ncog.RedisQueueWorker(Model(), redis_host=sys.argv[1], redis_port=sys.argv[2], input_queue=sys.argv[3], upload_url=sys.argv[4], consumer_id=sys.argv[5]).start()' > /usr/bin/cog-redis-queue-worker
 RUN chmod +x /usr/bin/cog-redis-queue-worker`
 }
diff --git a/pkg/serving/test.go b/pkg/serving/test.go
--- a/pkg/serving/test.go
+++ b/pkg/serving/test.go
@@ -75,7 +75,7 @@ func TestVersion(ctx context.Context, servingPlatform Platform, imageTag string,
 
 		input := NewExampleWithBaseDir(example.Input, dir)
 
-		result, err := deployment.RunInference(ctx, input, logWriter)
+		result, err := deployment.RunPrediction(ctx, input, logWriter)
 		if err != nil {
 			return nil, err
 		}
@@ -87,7 +87,7 @@ func TestVersion(ctx context.Context, servingPlatform Platform, imageTag string,
 		if err != nil {
 			return nil, fmt.Errorf("Failed to read output: %w", err)
 		}
-		logWriter.Infof("Inference result length: %d, mime type: %s", len(outputBytes), output.MimeType)
+		logWriter.Infof("Prediction length: %d, mime type: %s", len(outputBytes), output.MimeType)
 		if expectedOutput == nil {
 			updateExampleOutput(example, newExampleOutputs, outputBytes, output.MimeType, index)
 		} else {