adaptive-sdk 0.1.2__py3-none-any.whl → 0.1.3__py3-none-any.whl

adaptive_sdk/resources/models.py

@@ -12,6 +12,7 @@ from adaptive_sdk.graphql_client import (
     AttachModel,
     UpdateModelService,
     ModelData,
+    JobData,
     ModelServiceData,
     ListModelsModels,
     AddHFModelInput,
@@ -29,7 +30,9 @@ if TYPE_CHECKING:
 provider_config = {
     "open_ai": {
         "provider_data": lambda api_key, model_id: ModelProviderDataInput(
-            openAI=OpenAIProviderDataInput(apiKey=api_key, externalModelId=OpenAIModel(model_id))
+            openAI=OpenAIProviderDataInput(
+                apiKey=api_key, externalModelId=OpenAIModel(model_id)
+            )
         ),
     },
     "google": {
@@ -39,40 +42,53 @@ provider_config = {
     },
     "azure": {
         "provider_data": lambda api_key, model_id, endpoint: ModelProviderDataInput(
-            azure=AzureProviderDataInput(apiKey=api_key, externalModelId=model_id, endpoint=endpoint)
+            azure=AzureProviderDataInput(
+                apiKey=api_key, externalModelId=model_id, endpoint=endpoint
+            )
         )
     },
 }
 
 SupportedHFModels = Literal[
-    "google/gemma-3-4b-it",
-    "google/gemma-3-12b-it",
-    "google/gemma-3-27b-it",
+    "deepseek-ai/deepseek-coder-1.3b-base",
+    "deepseek-ai/deepseek-coder-6.7b-base",
+    "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B",
+    "deepseek-ai/DeepSeek-R1-Distill-Llama-8B",
+    "tiiuae/falcon-7b",
+    "tiiuae/falcon-7b-instruct",
+    "tiiuae/falcon-40b",
+    "tiiuae/falcon-180B",
+    "BAAI/bge-multilingual-gemma2",
+    "Locutusque/TinyMistral-248M",
+    "mistralai/Mistral-Small-24B-Instruct-2501",
+    "baffo32/decapoda-research-llama-7B-hf",
+    "princeton-nlp/Sheared-LLaMA-1.3B",
+    "meta-llama/Llama-3.1-8B",
     "meta-llama/Llama-3.1-8B-Instruct",
     "meta-llama/Llama-3.1-70B-Instruct",
-    "meta-llama/Llama-3.2-1B-Instruct",
-    "meta-llama/Llama-3.2-3B-Instruct",
     "meta-llama/Llama-3.3-70B-Instruct",
-    "mistralai/Mistral-Small-24B-Instruct-2501",
+    "nvidia/Llama3-ChatQA-1.5-70B",
+    "Qwen/Qwen2.5-0.5B",
     "Qwen/Qwen2.5-0.5B-Instruct",
-    "Qwen/Qwen2.5-1.5B-Instruct",
-    "Qwen/Qwen2.5-3B-Instruct",
-    "Qwen/Qwen2.5-7B-Instruct",
-    "Qwen/Qwen2.5-14B-Instruct",
-    "Qwen/Qwen2.5-32B-Instruct",
-    "Qwen/Qwen2.5-72B-Instruct",
-    "Qwen/Qwen2.5-Coder-0.5B-Instruct",
-    "Qwen/Qwen2.5-Coder-1.5B-Instruct",
-    "Qwen/Qwen2.5-Coder-3B-Instruct",
+    "Qwen/Qwen2.5-Coder-7B",
     "Qwen/Qwen2.5-Coder-7B-Instruct",
+    "Qwen/Qwen2.5-Math-7B",
+    "Qwen/Qwen2.5-Math-7B-Instruct",
     "Qwen/Qwen2.5-Coder-14B-Instruct",
     "Qwen/Qwen2.5-Coder-32B-Instruct",
+    "Qwen/QwQ-32B",
+    "google/gemma-3-1b-it",
+    "google/gemma-3-4b-it",
+    "google/gemma-3-12b-it",
+    "google/gemma-3-27b-it",
     "Qwen/Qwen3-0.6B",
     "Qwen/Qwen3-1.7B",
     "Qwen/Qwen3-4B",
     "Qwen/Qwen3-8B",
     "Qwen/Qwen3-14B",
     "Qwen/Qwen3-32B",
+    "01-ai/Yi-34B",
+    "HuggingFaceH4/zephyr-7b-beta",
 ]
 
 
@@ -80,7 +96,9 @@ def is_supported_model(model_id: str):
     supported_models = get_args(SupportedHFModels)
     if model_id not in supported_models:
         supported_models_str = "\n".join(supported_models)
-        raise ValueError(f"Model {model_id} is not supported.\n\nChoose from:\n{supported_models_str}")
+        raise ValueError(
+            f"Model {model_id} is not supported.\n\nChoose from:\n{supported_models_str}"
+        )
 
 
 class Models(SyncAPIResource, UseCaseResource):  # type: ignore[misc]
@@ -99,7 +117,7 @@ class Models(SyncAPIResource, UseCaseResource):  # type: ignore[misc]
         output_model_key: str,
         hf_token: str,
         compute_pool: str | None = None,
-    ) -> str:
+    ) -> JobData:
         """
         Add model from the HuggingFace Model hub to Adaptive model registry.
         It will take several minutes for the model to be downloaded and converted to Adaptive format.
@@ -145,16 +163,22 @@ class Models(SyncAPIResource, UseCaseResource):  # type: ignore[misc]
                 provider_data = provider_data_fn(api_key, external_model_id)
             case "azure":
                 if not endpoint:
-                    raise ValueError("`endpoint` is required to connect Azure external model.")
+                    raise ValueError(
+                        "`endpoint` is required to connect Azure external model."
+                    )
                 provider_data = provider_data_fn(api_key, external_model_id, endpoint)
             case _:
                 raise ValueError(f"Provider {provider} is not supported")
 
         provider_enum = ExternalModelProviderName(provider.upper())
-        input = AddExternalModelInput(name=name, provider=provider_enum, providerData=provider_data)
+        input = AddExternalModelInput(
+            name=name, provider=provider_enum, providerData=provider_data
+        )
         return self._gql_client.add_external_model(input).add_external_model
 
-    def list(self, filter: input_types.ModelFilter | None = None) -> Sequence[ListModelsModels]:
+    def list(
+        self, filter: input_types.ModelFilter | None = None
+    ) -> Sequence[ListModelsModels]:
         """
         List all models in Adaptive model registry.
         """
@@ -192,7 +216,9 @@ class Models(SyncAPIResource, UseCaseResource):  # type: ignore[misc]
             useCase=self.use_case_key(use_case),
             attached=True,
             wait=wait,
-            placement=(ModelPlacementInput.model_validate(placement) if placement else None),
+            placement=(
+                ModelPlacementInput.model_validate(placement) if placement else None
+            ),
         )
         result = self._gql_client.attach_model_to_use_case(input).attach_model
         if make_default:
@@ -253,7 +279,9 @@ class Models(SyncAPIResource, UseCaseResource):  # type: ignore[misc]
             isDefault=is_default,
             attached=attached,
             desiredOnline=desired_online,
-            placement=(ModelPlacementInput.model_validate(placement) if placement else None),
+            placement=(
+                ModelPlacementInput.model_validate(placement) if placement else None
+            ),
         )
         return self._gql_client.update_model(input).update_model_service
 
@@ -276,7 +304,9 @@ class Models(SyncAPIResource, UseCaseResource):  # type: ignore[misc]
             force: If model is attached to several use cases, `force` must equal `True` in order
                 for the model to be terminated.
         """
-        return self._gql_client.terminate_model(id_or_key=model, force=force).terminate_model
+        return self._gql_client.terminate_model(
+            id_or_key=model, force=force
+        ).terminate_model
 
 
 class AsyncModels(AsyncAPIResource, UseCaseResource):  # type: ignore[misc]
@@ -295,7 +325,7 @@ class AsyncModels(AsyncAPIResource, UseCaseResource):  # type: ignore[misc]
         output_model_key: str,
         hf_token: str,
         compute_pool: str | None = None,
-    ):
+    ) -> JobData:
         """
         Add model from the HuggingFace Model hub to Adaptive model registry.
         It will take several minutes for the model to be downloaded and converted to Adaptive format.
@@ -341,17 +371,23 @@ class AsyncModels(AsyncAPIResource, UseCaseResource):  # type: ignore[misc]
                 provider_data = provider_data_fn(api_key, external_model_id)
             case "azure":
                 if not endpoint:
-                    raise ValueError("`endpoint` is required to connect Azure external model.")
+                    raise ValueError(
+                        "`endpoint` is required to connect Azure external model."
+                    )
                 provider_data = provider_data_fn(api_key, external_model_id, endpoint)
             case _:
                 raise ValueError(f"Provider {provider} is not supported")
 
         provider_enum = ExternalModelProviderName(provider.upper())
-        input = AddExternalModelInput(name=name, provider=provider_enum, providerData=provider_data)
+        input = AddExternalModelInput(
+            name=name, provider=provider_enum, providerData=provider_data
+        )
         result = await self._gql_client.add_external_model(input)
         return result.add_external_model
 
-    async def list(self, filter: input_types.ModelFilter | None = None) -> Sequence[ListModelsModels]:
+    async def list(
+        self, filter: input_types.ModelFilter | None = None
+    ) -> Sequence[ListModelsModels]:
         """
         List all models in Adaptive model registry.
         """
@@ -389,7 +425,9 @@ class AsyncModels(AsyncAPIResource, UseCaseResource):  # type: ignore[misc]
             useCase=self.use_case_key(use_case),
             attached=True,
             wait=wait,
-            placement=(ModelPlacementInput.model_validate(placement) if placement else None),
+            placement=(
+                ModelPlacementInput.model_validate(placement) if placement else None
+            ),
         )
         result = await self._gql_client.attach_model_to_use_case(input)
         result = result.attach_model
@@ -453,7 +491,9 @@ class AsyncModels(AsyncAPIResource, UseCaseResource):  # type: ignore[misc]
             isDefault=is_default,
             attached=attached,
             desiredOnline=desired_online,
-            placement=(ModelPlacementInput.model_validate(placement) if placement else None),
+            placement=(
+                ModelPlacementInput.model_validate(placement) if placement else None
+            ),
         )
         result = await self._gql_client.update_model(input)
         return result.update_model_service
@@ -466,7 +506,9 @@ class AsyncModels(AsyncAPIResource, UseCaseResource):  # type: ignore[misc]
             model: Model key.
             wait: If `True`, call block until model is in `Online` state.
         """
-        return (await self._gql_client.deploy_model(id_or_key=model, wait=wait)).deploy_model
+        return (
+            await self._gql_client.deploy_model(id_or_key=model, wait=wait)
+        ).deploy_model
 
     async def terminate(self, model: str, force: bool = False) -> str:
         """
@@ -477,4 +519,6 @@ class AsyncModels(AsyncAPIResource, UseCaseResource):  # type: ignore[misc]
             force: If model is attached to several use cases, `force` must equal `True` in order
                 for the model to be terminated.
         """
-        return (await self._gql_client.terminate_model(id_or_key=model, force=force)).terminate_model
+        return (
+            await self._gql_client.terminate_model(id_or_key=model, force=force)
+        ).terminate_model

adaptive_sdk/resources/recipes.py

@@ -1,4 +1,10 @@
 from __future__ import annotations
+import os
+import io
+import zipfile
+import mimetypes
+from contextlib import contextmanager
+from loguru import logger
 from hypothesis_jsonschema import from_schema
 from typing import TYPE_CHECKING, Sequence, Any
 from pathlib import Path
@@ -18,20 +24,6 @@ from adaptive_sdk.graphql_client import (
 
 if TYPE_CHECKING:
     from adaptive_sdk.client import Adaptive, AsyncAdaptive
-import mimetypes
-
-
-def _count_keys_recursively(data: Any) -> int:
-    """Recursively counts the total number of keys in dictionaries within the data."""
-    count = 0
-    if isinstance(data, dict):
-        count += len(data)
-        for value in data.values():
-            count += _count_keys_recursively(value)
-    elif isinstance(data, list):
-        for item in data:
-            count += _count_keys_recursively(item)
-    return count
 
 
 class Recipes(SyncAPIResource, UseCaseResource):  # type: ignore[misc]
@@ -49,52 +41,35 @@ class Recipes(SyncAPIResource, UseCaseResource):  # type: ignore[misc]
 
     def upload(
         self,
-        file_path: str,
+        path: str,
         recipe_key: str,
         name: str | None = None,
         description: str | None = None,
         labels: dict[str, str] | None = None,
         use_case: str | None = None,
     ) -> CustomRecipeData:
-        filename = Path(file_path).stem
+        """
+        Upload a recipe from either a single Python file or a directory (path).
+        If a directory is provided, it must contain a 'main.py' and will be zipped in-memory before upload.
+        """
+        inferred_name = name or recipe_key
         label_inputs = [LabelInput(key=k, value=v) for k, v in labels.items()] if labels else None
         input = CreateRecipeInput(
             key=recipe_key,
-            name=name or filename,
+            name=inferred_name,
             description=description,
             labels=label_inputs,
         )
-        content_type = mimetypes.guess_type(file_path)[0] or "application/octet-stream"
-        with open(file_path, "rb") as f:
-            file_upload = Upload(filename=filename, content=f, content_type=content_type)
+        with _upload_from_path(path) as file_upload:
             return self._gql_client.create_custom_recipe(
                 use_case=self.use_case_key(use_case), input=input, file=file_upload
             ).create_custom_recipe
 
-    def run(
-        self,
-        recipe_key: str,
-        num_gpus: int,
-        input_args: dict | None = None,
-        name: str | None = None,
-        use_case: str | None = None,
-        compute_pool: str | None = None,
-    ) -> JobData:
-        input = JobInput(
-            recipe=recipe_key,
-            useCase=self.use_case_key(use_case),
-            args=input_args or {},
-            name=name,
-            computePool=compute_pool,
-            numGpus=num_gpus,
-        )
-        return self._gql_client.create_job(input).create_job
-
     def get(
         self,
         recipe_key: str,
         use_case: str | None = None,
-    ) -> CustomRecipeData:
+    ) -> CustomRecipeData | None:
         return self._gql_client.get_custom_recipe(
             id_or_key=recipe_key, use_case=self.use_case_key(use_case)
         ).custom_recipe
@@ -102,7 +77,7 @@ class Recipes(SyncAPIResource, UseCaseResource):  # type: ignore[misc]
     def update(
         self,
         recipe_key: str,
-        file_path: str | None = None,
+        path: str | None = None,
         name: str | None = None,
         description: str | None = None,
         labels: Sequence[tuple[str, str]] | None = None,
@@ -115,19 +90,21 @@ class Recipes(SyncAPIResource, UseCaseResource):  # type: ignore[misc]
             labels=label_inputs,
         )
 
-        file_upload = None
-        if file_path:
-            filename = Path(file_path).stem
-            content_type = mimetypes.guess_type(file_path)[0] or "application/octet-stream"
-            with open(file_path, "rb") as f:
-                file_upload = Upload(filename=filename, content=f, content_type=content_type)
-
-        return self._gql_client.update_custom_recipe(
-            use_case=self.use_case_key(use_case),
-            id=recipe_key,
-            input=input,
-            file=file_upload,
-        ).update_custom_recipe
+        if path:
+            with _upload_from_path(path) as file_upload:
+                return self._gql_client.update_custom_recipe(
+                    use_case=self.use_case_key(use_case),
+                    id=recipe_key,
+                    input=input,
+                    file=file_upload,
+                ).update_custom_recipe
+        else:
+            return self._gql_client.update_custom_recipe(
+                use_case=self.use_case_key(use_case),
+                id=recipe_key,
+                input=input,
+                file=None,
+            ).update_custom_recipe
 
     def delete(
         self,
@@ -140,6 +117,8 @@ class Recipes(SyncAPIResource, UseCaseResource):  # type: ignore[misc]
 
     def generate_sample_input(self, recipe_key: str, use_case: str | None = None) -> dict:
         recipe_details = self.get(recipe_key=recipe_key, use_case=self.use_case_key(use_case))
+        if recipe_details is None:
+            raise ValueError(f"Recipe {recipe_key} was not found")
         strategy = from_schema(recipe_details.json_schema)
 
         best_example = None
@@ -180,54 +159,33 @@ class AsyncRecipes(AsyncAPIResource, UseCaseResource):  # type: ignore[misc]
 
     async def upload(
         self,
-        file_path: str,
+        path: str,
         recipe_key: str,
         name: str | None = None,
         description: str | None = None,
         labels: Sequence[tuple[str, str]] | None = None,
         use_case: str | None = None,
     ) -> CustomRecipeData:
-        filename = Path(file_path).stem
+        inferred_name = name or recipe_key
         label_inputs = [LabelInput(key=k, value=v) for k, v in labels] if labels else None
         input = CreateRecipeInput(
             key=recipe_key,
-            name=name or filename,
+            name=inferred_name,
             description=description,
             labels=label_inputs,
         )
-        content_type = mimetypes.guess_type(file_path)[0] or "application/octet-stream"
-        with open(file_path, "rb") as f:
-            file_upload = Upload(filename=filename, content=f, content_type=content_type)
+        with _upload_from_path(path) as file_upload:
             return (
                 await self._gql_client.create_custom_recipe(
                     use_case=self.use_case_key(use_case), input=input, file=file_upload
                 )
             ).create_custom_recipe
 
-    async def run(
-        self,
-        recipe_key: str,
-        num_gpus: int,
-        input_args: dict | None = None,
-        name: str | None = None,
-        use_case: str | None = None,
-        compute_pool: str | None = None,
-    ) -> JobData:
-        input = JobInput(
-            recipe=recipe_key,
-            useCase=self.use_case_key(use_case),
-            args=input_args,
-            name=name,
-            computePool=compute_pool,
-            numGpus=num_gpus,
-        )
-        return (await self._gql_client.create_job(input)).create_job
-
     async def get(
         self,
         recipe_key: str,
         use_case: str | None = None,
-    ) -> CustomRecipeData:
+    ) -> CustomRecipeData | None:
         return (
             await self._gql_client.get_custom_recipe(id_or_key=recipe_key, use_case=self.use_case_key(use_case))
         ).custom_recipe
@@ -235,7 +193,7 @@ class AsyncRecipes(AsyncAPIResource, UseCaseResource):  # type: ignore[misc]
     async def update(
         self,
         recipe_key: str,
-        file_path: str | None = None,
+        path: str | None = None,
         name: str | None = None,
         description: str | None = None,
         labels: Sequence[tuple[str, str]] | None = None,
@@ -248,21 +206,25 @@ class AsyncRecipes(AsyncAPIResource, UseCaseResource):  # type: ignore[misc]
             labels=label_inputs,
         )
 
-        file_upload = None
-        if file_path:
-            filename = Path(file_path).stem
-            content_type = mimetypes.guess_type(file_path)[0] or "application/octet-stream"
-            with open(file_path, "rb") as f:
-                file_upload = Upload(filename=filename, content=f, content_type=content_type)
-
-        return (
-            await self._gql_client.update_custom_recipe(
-                use_case=self.use_case_key(use_case),
-                id=recipe_key,
-                input=input,
-                file=file_upload,
-            )
-        ).update_custom_recipe
+        if path:
+            with _upload_from_path(path) as file_upload:
+                return (
+                    await self._gql_client.update_custom_recipe(
+                        use_case=self.use_case_key(use_case),
+                        id=recipe_key,
+                        input=input,
+                        file=file_upload,
+                    )
+                ).update_custom_recipe
+        else:
+            return (
+                await self._gql_client.update_custom_recipe(
+                    use_case=self.use_case_key(use_case),
+                    id=recipe_key,
+                    input=input,
+                    file=None,
+                )
+            ).update_custom_recipe
 
     async def delete(
         self,
@@ -275,6 +237,8 @@ class AsyncRecipes(AsyncAPIResource, UseCaseResource):  # type: ignore[misc]
 
     async def generate_sample_input(self, recipe_key: str, use_case: str | None = None) -> dict:
         recipe_details = await self.get(recipe_key=recipe_key, use_case=self.use_case_key(use_case))
+        if recipe_details is None:
+            raise ValueError(f"Recipe {recipe_key} was not found")
         strategy = from_schema(recipe_details.json_schema)
 
         best_example = None
@@ -296,3 +260,102 @@ class AsyncRecipes(AsyncAPIResource, UseCaseResource):  # type: ignore[misc]
             print("A valid sample could not be generated. Returning an empty dict.")
             best_example = {}
         return dict(best_example)  # type: ignore
+
+
+def _count_keys_recursively(data: Any) -> int:
+    """Recursively counts the total number of keys in dictionaries within the data."""
+    count = 0
+    if isinstance(data, dict):
+        count += len(data)
+        for value in data.values():
+            count += _count_keys_recursively(value)
+    elif isinstance(data, list):
+        for item in data:
+            count += _count_keys_recursively(item)
+    return count
+
+
+def _validate_python_file(path: Path) -> None:
+    """Validate that the path exists, is a file and has a .py extension."""
+    if not path.exists():
+        raise FileNotFoundError(f"Python file not found: {path}")
+    if not path.is_file():
+        raise ValueError(f"Expected a file path, got a directory or non-file: {path}")
+    if path.suffix.lower() != ".py":
+        raise ValueError(f"Expected a Python file with .py extension, got: {path}")
+
+
+def _validate_recipe_directory(dir_path: Path) -> None:
+    """Validate that the directory exists and contains a main.py file."""
+    if not dir_path.exists():
+        raise FileNotFoundError(f"Directory not found: {dir_path}")
+    if not dir_path.is_dir():
+        raise ValueError(f"Expected a directory path, got a file: {dir_path}")
+    main_py = dir_path / "main.py"
+    if not main_py.exists() or not main_py.is_file():
+        raise FileNotFoundError(f"Directory must contain a 'main.py' file: {dir_path}")
+
+
+def _zip_directory_to_bytes_io(dir_path: Path) -> io.BytesIO:
+    """Zip the contents of a directory into an in-memory BytesIO buffer."""
+    buffer = io.BytesIO()
+    with zipfile.ZipFile(buffer, mode="w", compression=zipfile.ZIP_DEFLATED) as zf:
+        for root, _, files in os.walk(dir_path):
+            for file_name in files:
+                file_path = Path(root) / file_name
+                arcname = file_path.relative_to(dir_path)
+                zf.write(file_path, arcname.as_posix())
+    buffer.seek(0)
+    return buffer
+
+
+@contextmanager
+def _upload_from_path(path: str):
+    """
+    Context manager yielding an Upload object for a Python file or a directory.
+
+    - If path is a .py file, validates and opens it for upload.
+    - If path is a directory, validates it contains main.py, zips contents in-memory.
+    """
+    p = Path(path)
+    if p.is_file():
+        _validate_python_file(p)
+        filename = p.name
+        content_type = mimetypes.guess_type(str(p))[0] or "application/octet-stream"
+        f = open(p, "rb")
+        try:
+            yield Upload(filename=filename, content=f, content_type=content_type)
+        finally:
+            f.close()
+    elif p.is_dir():
+        _validate_recipe_directory(p)
+        # Ensure __init__.py exists at the root of the directory before zipping
+        created_init = False
+        root_init = p / "__init__.py"
+        zip_buffer = None
+        try:
+            if not root_init.exists():
+                root_init.touch()
+                created_init = True
+                logger.info(f"Added __init__.py to your directory, as it is required for proper execution of recipe")
+            zip_buffer = _zip_directory_to_bytes_io(p)
+        finally:
+            if created_init:
+                try:
+                    root_init.unlink()
+                    logger.info(f"Cleaned up __init__.py from your directory")
+                except Exception:
+                    logger.error(f"Failed to remove __init__.py from your directory")
+                    pass
+        if zip_buffer is None:
+            raise RuntimeError("Failed to create in-memory zip for directory upload")
+
+        filename = f"{p.name}.zip"
+        try:
+            yield Upload(filename=filename, content=zip_buffer, content_type="application/zip")
+        finally:
+            zip_buffer.close()
+    else:
+        if not p.exists():
+            raise FileNotFoundError(f"Path not found: {path}")
+        raise ValueError(f"Path must be a Python file or a directory: {path}")