Handle ResourceNotExistsError instead of 404 (#875)

dstackai · Feb 6, 2024 · 39037ae · 39037ae
1 parent e6e7147
commit 39037ae
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 19 deletions.
diff --git a/src/dstack/_internal/cli/commands/gateway.py b/src/dstack/_internal/cli/commands/gateway.py
@@ -60,7 +60,7 @@ def _register(self):
 
     def _command(self, args: argparse.Namespace):
         super()._command(args)
-        # TODO handle 404 and other errors
+        # TODO handle errors
         args.subfunc(args)
 
     def _list(self, args: argparse.Namespace):

diff --git a/src/dstack/api/_public/repos.py b/src/dstack/api/_public/repos.py
@@ -2,10 +2,9 @@
 from typing import Optional, Union
 
 import giturlparse
-import requests
 from git import InvalidGitRepositoryError
 
-from dstack._internal.core.errors import ConfigurationError
+from dstack._internal.core.errors import ConfigurationError, ResourceNotExistsError
 from dstack._internal.core.models.repos import LocalRepo, RemoteRepo
 from dstack._internal.core.models.repos.base import Repo, RepoType
 from dstack._internal.core.services.configs import ConfigManager
@@ -97,10 +96,8 @@ def is_initialized(
         try:
             self._api_client.repos.get(self._project, repo.repo_id, include_creds=False)
             return True
-        except requests.exceptions.HTTPError as e:
-            if e.response.status_code == 404:
-                return False
-            raise
+        except ResourceNotExistsError:
+            return False
 
     def load(
         self,
@@ -137,12 +134,10 @@ def load(
             repo = load_repo(repo_config)
             try:
                 self._api_client.repos.get(self._project, repo.repo_id, include_creds=False)
-            except requests.HTTPError as e:
-                if "404" in e.args[0]:
-                    raise ConfigurationError(
-                        "The repo is not initialized. Run `dstack init` to initialize the repo."
-                    )
-                raise
+            except ResourceNotExistsError:
+                raise ConfigurationError(
+                    "The repo is not initialized. Run `dstack init` to initialize the repo."
+                )
         else:
             logger.debug("Initializing repo")
             repo = LocalRepo(repo_dir=repo_dir)  # default

diff --git a/src/dstack/api/_public/runs.py b/src/dstack/api/_public/runs.py
@@ -9,11 +9,10 @@
 from pathlib import Path
 from typing import Dict, Iterable, List, Optional, Union
 
-import requests
 from websocket import WebSocketApp
 
 import dstack.api as api
-from dstack._internal.core.errors import ConfigurationError
+from dstack._internal.core.errors import ConfigurationError, ResourceNotExistsError
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.configurations import AnyRunConfiguration
 from dstack._internal.core.models.profiles import Profile, ProfileRetryPolicy, SpotPolicy
@@ -477,10 +476,8 @@ def get(self, run_name: str) -> Optional[Run]:
         try:
             run = self._api_client.runs.get(self._project, run_name)
             return self._model_to_run(run)
-        except requests.exceptions.HTTPError as e:
-            if e.response.status_code != 404:
-                raise
-        return None
+        except ResourceNotExistsError:
+            return None
 
     def _model_to_run(self, run: RunModel) -> Run:
         return Run(