diff --git a/crates/tower-cmd/src/run.rs b/crates/tower-cmd/src/run.rs index 5d818769..aad608a7 100644 --- a/crates/tower-cmd/src/run.rs +++ b/crates/tower-cmd/src/run.rs @@ -51,6 +51,10 @@ pub fn run_cmd() -> Command { pub async fn do_run(config: Config, args: &ArgMatches, cmd: Option<(&str, &ArgMatches)>) { let res = get_run_parameters(args, cmd); + // We always expect there to be an environment due to the fact that there is a + // default value. + let env = args.get_one::("environment").unwrap(); + match res { Ok((local, path, params, app_name)) => { debug!( @@ -64,13 +68,9 @@ pub async fn do_run(config: Config, args: &ArgMatches, cmd: Option<(&str, &ArgMa if app_name.is_some() { output::die("Running apps by name locally is not supported yet."); } else { - do_run_local(config, path, params).await; + do_run_local(config, path, env, params).await; } } else { - // We always expect there to be an environmnt due to the fact that there is a - // default value. - let env = args.get_one::("environment").unwrap(); - do_run_remote(config, path, env, params, app_name).await; } } @@ -83,10 +83,7 @@ pub async fn do_run(config: Config, args: &ArgMatches, cmd: Option<(&str, &ArgMa /// do_run_local is the entrypoint for running an app locally. It will load the Towerfile, build /// the package, and launch the app. The relevant package is cleaned up after execution is /// complete. -async fn do_run_local(config: Config, path: PathBuf, mut params: HashMap) { - // There is always an implicit `local` environment when running in a local context. - let env = "local".to_string(); - +async fn do_run_local(config: Config, path: PathBuf, env: &str, mut params: HashMap) { let mut spinner = output::spinner("Setting up runtime environment..."); // Load all the secrets and catalogs from the server @@ -104,6 +101,18 @@ async fn do_run_local(config: Config, path: PathBuf, mut params: HashMap = AppLauncher::default(); if let Err(err) = launcher - .launch(Context::new(), sender, package, env, secrets, params, catalogs) + .launch(Context::new(), sender, package, env.to_string(), secrets, params, env_vars) .await { output::runtime_error(err); diff --git a/src/tower/_client.py b/src/tower/_client.py index 0d3a0b0b..daa9b7e4 100644 --- a/src/tower/_client.py +++ b/src/tower/_client.py @@ -5,6 +5,7 @@ from ._context import TowerContext from .exceptions import ( + AppNotFoundError, NotFoundException, UnauthorizedException, UnknownException, @@ -24,6 +25,7 @@ RunAppResponse, ) from .tower_api_client.models.error_model import ErrorModel +from .tower_api_client.errors import UnexpectedStatus # WAIT_TIMEOUT is the amount of time to wait between requests when polling the # Tower API. @@ -84,18 +86,24 @@ def run_app( parameters=run_params, ) - output: Optional[Union[ErrorModel, RunAppResponse]] = run_app_api.sync( - slug=slug, client=client, body=input_body - ) + try: + output: Optional[Union[ErrorModel, RunAppResponse]] = run_app_api.sync( + slug=slug, client=client, body=input_body + ) - if output is None: - raise RuntimeError("Error running app") - else: - if isinstance(output, ErrorModel): - raise RuntimeError(f"Error running app: {output.title}") + if output is None: + raise RuntimeError("Error running app") else: - return output.run - + if isinstance(output, ErrorModel): + raise RuntimeError(f"Error running app: {output.title}") + else: + return output.run + except UnexpectedStatus as e: + # Raise an AppNotFoundError here if the app was, indeed, not found. + if e.status_code == 404: + raise AppNotFoundError(slug) + else: + raise UnknownException(f"Unexpected status code {e.status_code} when running app {slug}") def wait_for_run( run: Run, @@ -303,13 +311,23 @@ def _env_client(ctx: TowerContext, timeout: Optional[float] = None) -> Authentic else: tower_url += "/v1" + if ctx.jwt is not None: + token = ctx.jwt + auth_header_name = "Authorization" + prefix = "Bearer" + else: + token = ctx.api_key + auth_header_name = "X-API-Key" + prefix = "" + return AuthenticatedClient( verify_ssl=False, base_url=tower_url, - token=ctx.api_key, - auth_header_name="X-API-Key", - prefix="", + token=token, + auth_header_name=auth_header_name, + prefix=prefix, timeout=timeout, + raise_on_unexpected_status=True, ) diff --git a/src/tower/_context.py b/src/tower/_context.py index c2c909fb..995dc9f2 100644 --- a/src/tower/_context.py +++ b/src/tower/_context.py @@ -1,10 +1,11 @@ import os class TowerContext: - def __init__(self, tower_url: str, environment: str, api_key: str = None, hugging_face_provider: str = None, hugging_face_api_key: str = None): + def __init__(self, tower_url: str, environment: str, api_key: str = None, hugging_face_provider: str = None, hugging_face_api_key: str = None, jwt: str = None): self.tower_url = tower_url self.environment = environment self.api_key = api_key + self.jwt = jwt self.hugging_face_provider = hugging_face_provider self.hugging_face_api_key = hugging_face_api_key @@ -18,9 +19,13 @@ def is_local(self) -> bool: @classmethod def build(cls): - tower_url = os.getenv("TOWER_URL") - tower_environment = os.getenv("TOWER_ENVIRONMENT") + tower_url = os.getenv("TOWER_URL", "https://api.tower.dev") + tower_environment = os.getenv("TOWER_ENVIRONMENT", "default") tower_api_key = os.getenv("TOWER_API_KEY") + tower_jwt = os.getenv("TOWER_JWT") + + # NOTE: These are experimental, used only for our experimental Hugging + # Face integration for LLMs. hugging_face_provider = os.getenv("TOWER_HUGGING_FACE_PROVIDER") hugging_face_api_key = os.getenv("TOWER_HUGGING_FACE_API_KEY") @@ -30,5 +35,6 @@ def build(cls): api_key = tower_api_key, hugging_face_provider = hugging_face_provider, hugging_face_api_key = hugging_face_api_key, + jwt = tower_jwt, ) diff --git a/src/tower/exceptions.py b/src/tower/exceptions.py index e2a7295f..cdd845e5 100644 --- a/src/tower/exceptions.py +++ b/src/tower/exceptions.py @@ -27,3 +27,7 @@ def __init__(self, time: float): class RunFailedError(RuntimeError): def __init__(self, app_name: str, number: int, state: str): super().__init__(f"Run {app_name}#{number} failed with status '{state}'") + +class AppNotFoundError(RuntimeError): + def __init__(self, app_name: str): + super().__init__(f"App '{app_name}' not found in the Tower.") diff --git a/tests/tower/test_client.py b/tests/tower/test_client.py index 13598c36..584e15a5 100644 --- a/tests/tower/test_client.py +++ b/tests/tower/test_client.py @@ -324,3 +324,67 @@ def test_raising_an_error_during_partial_failure( # Now actually wait for the runs with pytest.raises(RunFailedError) as excinfo: tower.wait_for_runs(runs, raise_on_failure=True) + + +def test_raising_an_error_for_a_not_found_app( + httpx_mock, + mock_api_config, + mock_run_response_factory, + create_run_object +): + tower = mock_api_config + + # Mock a 404 response with error model + error_response = { + "$schema": "https://api.tower.dev/v1/schemas/ErrorModel.json", + "status": 404, + "title": "Not Found", + "detail": "The requested app 'non-existent-app' was not found", + "instance": "https://api.example.com/v1/apps/non-existent-app", + "type": "about:blank" + } + + httpx_mock.add_response( + method="POST", + url="https://api.example.com/v1/apps/non-existent-app/runs", + json=error_response, + status_code=404, + ) + + # Attempt to run a non-existent app and verify it raises the correct error + with pytest.raises(tower.exceptions.AppNotFoundError) as excinfo: + tower.run_app("non-existent-app") + + assert "not found" in str(excinfo.value).lower() + + +def test_raising_an_unexpected_error_based_on_status_code( + httpx_mock, + mock_api_config, + mock_run_response_factory, + create_run_object +): + tower = mock_api_config + + # Mock a 404 response with error model + error_response = { + "$schema": "https://api.tower.dev/v1/schemas/ErrorModel.json", + "status": 404, + "title": "Not Found", + "detail": "The requested app 'non-existent-app' was not found", + "instance": "https://api.example.com/v1/apps/non-existent-app", + "type": "about:blank" + } + + httpx_mock.add_response( + method="POST", + url="https://api.example.com/v1/apps/non-existent-app/runs", + json=error_response, + status_code=400, + ) + + # Attempt to run a non-existent app and verify it raises the correct error + with pytest.raises(tower.exceptions.UnknownException) as excinfo: + tower.run_app("non-existent-app") + + assert "unexpected status code" in str(excinfo.value).lower()