Skip to content

Commit

Permalink
Merge branch 'master' into 2PR
Browse files Browse the repository at this point in the history
  • Loading branch information
samruds authored Mar 4, 2024
2 parents 9aa52b1 + 790bd87 commit 1cbb9ce
Show file tree
Hide file tree
Showing 3 changed files with 8 additions and 5 deletions.
4 changes: 3 additions & 1 deletion src/sagemaker/local/image.py
Original file line number Diff line number Diff line change
Expand Up @@ -860,7 +860,9 @@ def _create_docker_host(
# to setting --runtime=nvidia in the docker commandline.
if self.instance_type == "local_gpu":
host_config["deploy"] = {
"resources": {"reservations": {"devices": [{"capabilities": ["gpu"]}]}}
"resources": {
"reservations": {"devices": [{"count": "all", "capabilities": ["gpu"]}]}
}
}

if not self.is_studio and command == "serve":
Expand Down
2 changes: 1 addition & 1 deletion tests/unit/sagemaker/local/test_local_image.py
Original file line number Diff line number Diff line change
Expand Up @@ -871,7 +871,7 @@ def test_container_has_gpu_support(tmpdir, sagemaker_session):
docker_host = sagemaker_container._create_docker_host("host-1", {}, set(), "train", [])
assert "deploy" in docker_host
assert docker_host["deploy"] == {
"resources": {"reservations": {"devices": [{"capabilities": ["gpu"]}]}}
"resources": {"reservations": {"devices": [{"count": "all", "capabilities": ["gpu"]}]}}
}


Expand Down
7 changes: 4 additions & 3 deletions tox.ini
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
# and then run "tox" from this directory.

[tox]
envlist = black-format,flake8,pylint,docstyle,sphinx,doc8,twine,py38,py39,py310
envlist = black-format,flake8,pylint,docstyle,sphinx,doc8,twine,py38,py39,py310,py311

skip_missing_interpreters = False

Expand Down Expand Up @@ -84,12 +84,13 @@ commands =
pip install 'apache-airflow==2.8.2' --constraint "https://raw.githubusercontent.com/apache/airflow/constraints-2.8.2/constraints-3.8.txt"
pip install 'torch==2.0.1+cpu' -f 'https://download.pytorch.org/whl/torch_stable.html'
pip install 'torchvision==0.15.2+cpu' -f 'https://download.pytorch.org/whl/torch_stable.html'
pip install 'dill>=0.3.8'

pytest --cov=sagemaker --cov-append {posargs}
{env:IGNORE_COVERAGE:} coverage report -i --fail-under=86
{env:IGNORE_COVERAGE:} coverage report -i --fail-under=86
deps = .[test]
depends =
{py38,py39,py310}: clean
{py38,py39,py310,p311}: clean

[testenv:flake8]
skipdist = true
Expand Down

0 comments on commit 1cbb9ce

Please sign in to comment.