add a new flag that enables users to disable serving openapi.json and…

… docs swagger ui endpoints per github.com/replicate/issues/1986

pkg/cli/build.go

@@ -25,6 +25,7 @@ var buildUseCogBaseImage bool
 var buildStrip bool
 var buildPrecompile bool
 var buildFast bool
+var disableDocs bool
 
 const useCogBaseImageFlagKey = "use-cog-base-image"
 
@@ -48,6 +49,7 @@ func newBuildCommand() *cobra.Command {
 	addStripFlag(cmd)
 	addPrecompileFlag(cmd)
 	addFastFlag(cmd)
+	addDocsFlag(cmd)
 	cmd.Flags().StringVarP(&buildTag, "tag", "t", "", "A name for the built image in the form 'repository:tag'")
 	return cmd
 }
@@ -144,6 +146,12 @@ func addFastFlag(cmd *cobra.Command) {
 	_ = cmd.Flags().MarkHidden(fastFlag)
 }
 
+func addDocsFlag(cmd *cobra.Command) {
+	const docsFlag = "disable-docs"
+	cmd.Flags().BoolVar(&disableDocs, docsFlag, false, "Whether to disable Swagger UI for docs and openapi.json")
+	// _ = cmd.Flags().MarkHidden(docsFlag)
+}
+
 func checkMutuallyExclusiveFlags(cmd *cobra.Command, args []string) error {
 	flags := []string{useCogBaseImageFlagKey, "use-cuda-base-image", "dockerfile"}
 	var flagsSet []string

pkg/cli/serve.go

@@ -34,6 +34,7 @@ Generate and run an HTTP server based on the declared model inputs and outputs.`
 	addUseCogBaseImageFlag(cmd)
 	addGpusFlag(cmd)
 	addFastFlag(cmd)
+	addDocsFlag(cmd)
 
 	cmd.Flags().IntVarP(&port, "port", "p", port, "Port on which to listen")
 
@@ -69,6 +70,10 @@ func cmdServe(cmd *cobra.Command, arg []string) error {
 		"--await-explicit-shutdown", "true",
 	}
 
+	if disableDocs {
+		args = append(args, "--disable-docs")
+	}
+
 	runOptions := docker.RunOptions{
 		Args:    args,
 		Env:     envFlags,

python/cog/server/http.py

@@ -117,10 +117,15 @@ def create_app(  # pylint: disable=too-many-arguments,too-many-locals,too-many-s
     mode: Mode = Mode.PREDICT,
     is_build: bool = False,
     await_explicit_shutdown: bool = False,  # pylint: disable=redefined-outer-name
+    enable_docs: bool = True,
 ) -> MyFastAPI:
+
+    openapi_url = "/openapi.json" if enable_docs else None
+
     app = MyFastAPI(  # pylint: disable=redefined-outer-name
         title="Cog",  # TODO: mention model name?
         # version=None # TODO
+        openapi_url=openapi_url
     )
 
     def custom_openapi() -> Dict[str, Any]:
@@ -141,7 +146,8 @@ def custom_openapi() -> Dict[str, Any]:
 
         return app.openapi_schema
 
-    app.openapi = custom_openapi
+    if enable_docs:
+        app.openapi = custom_openapi
 
     app.state.health = Health.STARTING
     app.state.setup_result = None
@@ -192,8 +198,8 @@ async def wrapped(*args: "P.args", **kwargs: "P.kwargs") -> "T":  # pylint: disa
 
     index_document = {
         "cog_version": __version__,
-        "docs_url": "/docs",
-        "openapi_url": "/openapi.json",
+        "docs_url": "/docs" if enable_docs else None,
+        "openapi_url": "/openapi.json" if enable_docs else None,
         "shutdown_url": "/shutdown",
         "healthcheck_url": "/health-check",
         "predictions_url": "/predictions",
@@ -626,7 +632,16 @@ def _cpu_count() -> int:
         choices=list(Mode),
         help="Experimental: Run in 'predict' or 'train' mode",
     )
+    parser.add_argument(
+        "--disable-docs",
+        action="store_true", 
+        dest="disable_docs",
+        default=False, 
+        help="Disable OpenAPI schema and Swagger UI for docs"
+    )
     args = parser.parse_args()
+    # DO NOT SUBMIT: Print the parsed arguments for debugging only
+    print(f"### Parsed arguments: {args}") 
 
     if args.version:
         print(f"cog.server.http {__version__}")
@@ -654,6 +669,7 @@ def _cpu_count() -> int:
         upload_url=args.upload_url,
         mode=args.mode,
         await_explicit_shutdown=await_explicit_shutdown,
+        enable_docs=not args.disable_docs,
     )
 
     host: str = args.host